Coverage for mindsdb / integrations / handlers / clipdrop_handler / clipdrop.py: 0%

1import requests 

2import uuid 

3 

4 

5class ClipdropClient: 

6 def __init__(self, api_key, local_dir): 

7 self.api_key = api_key 

8 self.dir_to_save = local_dir if local_dir.endswith("/") else local_dir + "/" 

9 self.base_endpoint = "https://clipdrop-api.co/" 

10 

11 def write_to_file(self, img): 

12 f_name = self.dir_to_save + str(uuid.uuid4()) + ".png" 

13 with open(f_name, "wb") as f: 

14 f.write(img) 

15 return f_name 

16 

17 def make_request(self, url, files={}, data={}): 

18 headers = {"x-api-key": self.api_key} 

19 resp = requests.post(url, headers=headers, files=files, data=data) 

20 if resp.status_code == 200: 

21 saved_path = self.write_to_file(resp.content) 

22 return saved_path 

23 return resp.text 

24 

25 def image_extension_check(self, url): 

26 ext = url.split(".")[-1] 

27 if ext in ["jpeg", "jpg", "png"]: 

28 return ext 

29 raise Exception("Unknown image format. Currently jpg, jpeg & png are supported.") 

30 

31 def download_image(self, path): 

32 

33 img_ext = self.image_extension_check(path) 

34 try: 

35 res = requests.get(path) 

36 return {"img_ext": img_ext, "content": res.content} 

37 except Exception as e: 

38 raise Exception(f"Failed to download image: {e}") 

39 

40 def remove_text(self, img_url): 

41 url = f'{self.base_endpoint}remove-text/v1' 

42 img_content = self.download_image(img_url) 

43 files = { 

44 'image_file': ('image.jpg', img_content["content"], f'image/{img_content["img_ext"]}') 

45 } 

46 return self.make_request(url, files=files) 

47 

48 def remove_background(self, img_url): 

49 url = f'{self.base_endpoint}remove-background/v1' 

50 img_content = self.download_image(img_url) 

51 files = { 

52 'image_file': ('image.jpg', img_content["content"], f'image/{img_content["img_ext"]}') 

53 } 

54 return self.make_request(url, files=files) 

55 

56 def sketch_to_image(self, img_url, prompt): 

57 url = f'{self.base_endpoint}sketch-to-image/v1/sketch-to-image' 

58 img_content = self.download_image(img_url) 

59 files = { 

60 'sketch_file': ('image.jpg', img_content["content"], f'image/{img_content["img_ext"]}') 

61 } 

62 data = {'prompt': prompt} 

63 return self.make_request(url, files=files, data=data) 

64 

65 def text_to_image(self, prompt): 

66 url = f'{self.base_endpoint}text-to-image/v1' 

67 files = { 

68 'prompt': (None, prompt, 'text/plain') 

69 } 

70 return self.make_request(url, files=files) 

71 

72 def replace_background(self, img_url, prompt): 

73 url = f'{self.base_endpoint}replace-background/v1' 

74 img_content = self.download_image(img_url) 

75 files = { 

76 'image_file': ('image.jpg', img_content["content"], f'image/{img_content["img_ext"]}') 

77 } 

78 data = {'prompt': prompt} 

79 return self.make_request(url, files=files, data=data) 

80 

81 def reimagine(self, img_url): 

82 url = f'{self.base_endpoint}reimagine/v1/reimagine' 

83 img_content = self.download_image(img_url) 

84 files = { 

85 'image_file': ('image.jpg', img_content["content"], f'image/{img_content["img_ext"]}') 

86 } 

87 return self.make_request(url, files=files)