Create code.py

code.py

@@ -0,0 +1,90 @@
+import torch
+from transformers import AutoTokenizer, VisionEncoderDecoderModel, AutoImageProcessor
+from PIL import Image
+from torchvision.transforms.functional import crop
+import gradio as gr
+
+# Load models during initialization
+def init():
+    global object_detection_model, captioning_model, tokenizer, captioning_processor
+
+    # Step 1: Load the YOLOv5 model from Hugging Face
+    try:
+        object_detection_model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5/weights/best14.pt')  # Assuming model is locally available
+        print("YOLOv5 model loaded successfully.")
+    except Exception as e:
+        print(f"Error loading YOLOv5 model: {e}")
+
+    # Step 2: Load the ViT-GPT2 captioning model from Hugging Face
+    try:
+        captioning_model = VisionEncoderDecoderModel.from_pretrained("motheecreator/ViT-GPT2-Image-Captioning")
+        tokenizer = AutoTokenizer.from_pretrained("motheecreator/ViT-GPT2-Image-Captioning")
+        captioning_processor = AutoImageProcessor.from_pretrained("motheecreator/ViT-GPT2-Image-Captioning")
+        print("ViT-GPT2 model loaded successfully.")
+    except Exception as e:
+        print(f"Error loading captioning model: {e}")
+
+# Utility function to crop objects from the image based on bounding boxes
+def crop_objects(image, boxes):
+    cropped_images = []
+    for box in boxes:
+        cropped_image = crop(image, int(box[1]), int(box[0]), int(box[3] - box[1]), int(box[2] - box[0]))
+        cropped_images.append(cropped_image)
+    return cropped_images
+
+# Gradio interface function
+def process_image(image):
+    try:
+        # Step 1: Perform object detection with YOLOv5
+        results = object_detection_model(image)
+        boxes = results.xyxy[0][:, :4].cpu().numpy()  # Bounding boxes
+        labels = [results.names[int(class_id)] for class_id in results.xyxy[0][:, 5].cpu().numpy().astype(int)]  # Class names
+        scores = results.xyxy[0][:, 4].cpu().numpy()  # Confidence scores
+
+        # Step 2: Generate caption for the whole image
+        original_inputs = captioning_processor(images=image, return_tensors="pt")
+        with torch.no_grad():
+            original_caption_ids = captioning_model.generate(**original_inputs)
+        original_caption = tokenizer.decode(original_caption_ids[0], skip_special_tokens=True)
+
+        # Step 3: Crop detected objects and generate captions for each object
+        cropped_images = crop_objects(image, boxes)
+        captions = []
+        for cropped_image in cropped_images:
+            inputs = captioning_processor(images=cropped_image, return_tensors="pt")
+            with torch.no_grad():
+                caption_ids = captioning_model.generate(**inputs)
+            caption = tokenizer.decode(caption_ids[0], skip_special_tokens=True)
+            captions.append(caption)
+
+        # Prepare the result for visualization
+        detection_results = []
+        for i, (label, box, score, caption) in enumerate(zip(labels, boxes, scores, captions)):
+            detection_results.append({
+                "label": label,
+                "caption": caption,
+                "bounding_box": [float(coord) for coord in box],  # Convert to float
+                "confidence_score": float(score)  # Convert to float
+            })
+
+        # Return the image with detections and the caption
+        return results.render()[0], detection_results, original_caption
+
+    except Exception as e:
+        return None, {"error": str(e)}, None
+
+# Initialize models
+init()
+
+# Gradio Interface
+interface = gr.Interface(
+    fn=process_image,  # Function to run
+    inputs=gr.inputs.Image(type="pil"),  # Input: Image upload
+    outputs=[gr.outputs.Image(type="pil", label="Detected Objects"),  # Output 1: Image with bounding boxes
+             gr.outputs.JSON(label="Object Captions & Bounding Boxes"),  # Output 2: JSON results for each object
+             gr.outputs.Textbox(label="Whole Image Caption")],  # Output 3: Caption for the whole image
+    live=True
+)
+
+# Launch the Gradio app
+interface.launch()