Reverb/Embrace-Vision
0
1import gradio as gr2from transformers import CLIPProcessor, CLIPModel, CLIPTokenizer, BlipProcessor, BlipForConditionalGeneration, pipeline3from sentence_transformers import SentenceTransformer, util4import pickle5import numpy as np6from PIL import Image, ImageEnhance7import os8import io9import concurrent.futures10import warnings11 12warnings.filterwarnings("ignore")13 14class CLIPModelHandler:15 def __init__(self, model_name):16 self.model_name = model_name17 self.img_names, self.img_emb = self.load_precomputed_embeddings()18 19 def load_precomputed_embeddings(self):20 emb_filename = 'unsplash-25k-photos-embeddings.pkl'21 with open(emb_filename, 'rb') as fIn:22 img_names, img_emb = pickle.load(fIn)23 return img_names, img_emb24 25 def search_text(self, query, top_k=1):26 model = CLIPModel.from_pretrained(self.model_name)27 processor = CLIPProcessor.from_pretrained(self.model_name)28 tokenizer = CLIPTokenizer.from_pretrained(self.model_name)29 30 inputs = tokenizer([query], padding=True, return_tensors="pt")31 query_emb = model.get_text_features(**inputs)32 hits = util.semantic_search(query_emb, self.img_emb, top_k=top_k)[0]33 34 images = [Image.open(os.path.join("photos/", self.img_names[hit['corpus_id']])) for hit in hits]35 return images36 37 def search_image(self, image_path, top_k=1):38 model = CLIPModel.from_pretrained(self.model_name)39 processor = CLIPProcessor.from_pretrained(self.model_name)40 41 # Load and preprocess the image42 image = Image.open(image_path)43 inputs = processor(images=image, return_tensors="pt")44 45 # Get the image features46 outputs = model(**inputs)47 image_emb = outputs.logits_per_image48 49 # Perform semantic search50 hits = util.semantic_search(image_emb, self.img_emb, top_k=top_k)[0]51 52 # Retrieve and return the relevant images53 result_images = []54 for hit in hits:55 img = Image.open(os.path.join("photos/", self.img_names[hit['corpus_id']]))56 result_images.append(img)57 58 return result_images59 60class BLIPImageCaptioning:61 def __init__(self, blip_model_name):62 self.blip_model_name = blip_model_name63 64 def preprocess_image(self, image):65 if isinstance(image, str):66 return Image.open(image).convert('RGB')67 elif isinstance(image, np.ndarray):68 return Image.fromarray(np.uint8(image)).convert('RGB')69 else:70 raise ValueError("Invalid input type for image. Supported types: str (file path) or np.ndarray.")71 72 def generate_caption(self, image):73 try:74 model = BlipForConditionalGeneration.from_pretrained(self.blip_model_name)75 processor = BlipProcessor.from_pretrained(self.blip_model_name)76 77 raw_image = self.preprocess_image(image)78 inputs = processor(raw_image, return_tensors="pt")79 out = model.generate(**inputs)80 unconditional_caption = processor.decode(out[0], skip_special_tokens=True)81 82 return unconditional_caption83 except Exception as e:84 return {"error": str(e)}85 86 def generate_captions_parallel(self, images):87 with concurrent.futures.ThreadPoolExecutor() as executor:88 results = list(executor.map(self.generate_caption, images))89 90 return results91 92 93# Initialize the CLIP model handler94clip_handler = CLIPModelHandler("openai/clip-vit-base-patch32")95 96# Initialize the zero-shot image classification pipeline97clip_classifier = pipeline("zero-shot-image-classification", model="openai/clip-vit-base-patch32")98 99# Load BLIP model directly100blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")101blip_model_name = "Salesforce/blip-image-captioning-base"102 103# Function for text-to-image search104def text_to_image_interface(query, top_k):105 try:106 # Perform text-to-image search107 result_images = clip_handler.search_text(query, top_k)108 109 # Resize images before displaying110 result_images_resized = [image.resize((224, 224)) for image in result_images]111 112 # Display more information about the results113 result_info = [{"Image Name": os.path.basename(img_path)} for img_path in clip_handler.img_names]114 115 return result_images_resized, result_info116 except Exception as e:117 return gr.Error(f"Error in text-to-image search: {str(e)}")118 119 120# Gradio Interface function for zero-shot classification121def zero_shot_classification(image, labels_text):122 try:123 # Convert image to PIL format124 PIL_image = Image.fromarray(np.uint8(image)).convert('RGB')125 126 # Split labels_text into a list of labels127 labels = labels_text.split(",")128 129 # Perform zero-shot classification130 res = clip_classifier(images=PIL_image, candidate_labels=labels, hypothesis_template="This is a photo of a {}")131 132 # Format the result as a dictionary133 formatted_results = {dic["label"]: dic["score"] for dic in res}134 135 return formatted_results136 except Exception as e:137 return gr.Error(f"Error in zero-shot classification: {str(e)}")138 139 140 141def preprocessing_interface(original_image, brightness_slider, contrast_slider, saturation_slider, sharpness_slider, rotation_slider):142 try:143 # Convert NumPy array to PIL Image144 PIL_image = Image.fromarray(np.uint8(original_image)).convert('RGB')145 146 # Normalize slider values to be in the range [0, 1]147 brightness_normalized = brightness_slider / 100.0148 contrast_normalized = contrast_slider / 100.0149 saturation_normalized = saturation_slider / 100.0150 sharpness_normalized = sharpness_slider / 100.0151 152 # Apply preprocessing based on user input153 PIL_image = PIL_image.convert("RGB")154 PIL_image = PIL_image.rotate(rotation_slider)155 156 # Adjust brightness157 enhancer = ImageEnhance.Brightness(PIL_image)158 PIL_image = enhancer.enhance(brightness_normalized)159 160 # Adjust contrast161 enhancer = ImageEnhance.Contrast(PIL_image)162 PIL_image = enhancer.enhance(contrast_normalized)163 164 # Adjust saturation165 enhancer = ImageEnhance.Color(PIL_image)166 PIL_image = enhancer.enhance(saturation_normalized)167 168 # Adjust sharpness169 enhancer = ImageEnhance.Sharpness(PIL_image)170 PIL_image = enhancer.enhance(sharpness_normalized)171 172 # Return the processed image173 return PIL_image174 except Exception as e:175 return gr.Error(f"Error in preprocessing: {str(e)}")176 177def generate_captions(images):178 blip_model = BlipForConditionalGeneration.from_pretrained(blip_model_name)179 blip_processor = BlipProcessor.from_pretrained(blip_model_name)180 181 return [blip_model_instance.generate_caption(image) for image in images]182 183 184# Gradio Interfaces185zero_shot_classification_interface = gr.Interface(186 fn=zero_shot_classification,187 inputs=[188 gr.Image(label="Image Query", elem_id="image_input"),189 gr.Textbox(label="Labels (comma-separated)", elem_id="labels_input"),190 ],191 outputs=gr.Label(elem_id="label_image"),192)193 194text_to_image_interface = gr.Interface(195 fn=text_to_image_interface,196 inputs=[197 gr.Textbox(198 lines=2,199 label="Text Query",200 placeholder="Enter text here...",201 ),202 gr.Slider(0, 5, step=1, label="Top K Results"),203 ],204 outputs=[205 gr.Gallery(206 label="Text-to-Image Search Results",207 elem_id="gallery_text",208 grid_cols=2,209 height="auto",210 ),211 gr.Text(label="Result Information", elem_id="text_info"),212 ],213)214 215blip_model = BLIPImageCaptioning(blip_model_name) # Instantiate the object216blip_captioning_interface = gr.Interface(217 fn=blip_model.generate_caption, # Correct the method name218 inputs=gr.Image(label="Image for Captioning", elem_id="blip_caption_image"),219 outputs=gr.Textbox(label="Generated Captions", elem_id="blip_generated_captions", default=""),220)221 222preprocessing_interface = gr.Interface(223 fn=blip_model.preprocess_image, # Correct the method name224 inputs=[225 gr.Image(label="Original Image", elem_id="original_image"),226 ],227 outputs=[228 gr.Image(label="Processed Image", elem_id="processed_image"),229 ],230)231 232# Tabbed Interface233app = gr.TabbedInterface(234 interface_list=[text_to_image_interface, zero_shot_classification_interface, blip_captioning_interface],235 tab_names=["Text-to-Image Search", "Zero-Shot Classification", "BLIP Image Captioning"],236)237 238# Launch the Gradio interface239app.launch(debug=True, share="true")