Update app.py

app.py

@@ -1,68 +1,69 @@
 import gradio as gr
 import torch
-import numpy as np
 from PIL import Image
-import matplotlib.pyplot as plt
-import io
-from sam2.sam2_image_predictor import SAM2ImagePredictor
-
-# Set up device
-device = "cuda" if torch.cuda.is_available() else "cpu"
+import cv2
+import numpy as np
+from transformers import CLIPProcessor, CLIPModel
+from ultralytics import FastSAM
+from ultralytics.models.fastsam import FastSAMPrompt
 
-# Load SAM 2 model
-predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")
+# Load CLIP model
+model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
-def fig2img(fig):
-    buf = io.BytesIO()
-    fig.savefig(buf)
-    buf.seek(0)
-    img = Image.open(buf)
-    return img
+# Load FastSAM model
+fast_sam = FastSAM('FastSAM-x.pt')
 
-def plot_masks(image, masks):
-    fig, ax = plt.subplots(figsize=(10, 10))
-    ax.imshow(image)
+def process_image_clip(image, text_input):
+    # Process image for CLIP
+    inputs = processor(
+        images=image,
+        text=[text_input],
+        return_tensors="pt",
+        padding=True
+    )
+    
+    # Get model predictions
+    outputs = model(**inputs)
+    logits_per_image = outputs.logits_per_image
+    probs = logits_per_image.softmax(dim=1)
     
-    for mask in masks:
-        masked = np.ma.masked_where(mask == 0, mask)
-        ax.imshow(masked, alpha=0.5, cmap=plt.cm.get_cmap('jet'))
-    ax.axis('off')
-    plt.close()
-    return fig2img(fig)
+    confidence = float(probs[0][0])
+    return f"Confidence that the image contains '{text_input}': {confidence:.2%}"
 
-def segment_everything(input_image):
-    try:
-        if input_image is None:
-            return None, "Please upload an image before submitting."
-        
-        input_image = Image.fromarray(input_image).convert("RGB")
-        
-        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
-            predictor.set_image(input_image)
-            # Use 'everything' prompt
-            masks, _, _ = predictor.predict([])
-        
-        # Plot the results
-        result_image = plot_masks(input_image, masks)
-        
-        return result_image, f"Segmented everything in the image. Found {len(masks)} objects."
+def process_image_fastsam(image):
+    # Convert PIL image to numpy array
+    image_np = np.array(image)
     
-    except Exception as e:
-        return None, f"An error occurred: {str(e)}"
+    # Run FastSAM inference
+    everything_results = fast_sam(image_np, device='cpu', retina_masks=True, imgsz=1024, conf=0.4, iou=0.9)
+    prompt_process = FastSAMPrompt(image_np, everything_results, device='cpu')
+    
+    # Get everything mask
+    ann = prompt_process.everything()
+    
+    # Convert annotation to image
+    result_image = prompt_process.plot_to_result()
+    return Image.fromarray(result_image)
 
 # Create Gradio interface
-iface = gr.Interface(
-    fn=segment_everything,
-    inputs=[
-        gr.Image(type="numpy", label="Upload an image")
-    ],
-    outputs=[
-        gr.Image(type="pil", label="Segmented Image"),
-        gr.Textbox(label="Status")
-    ],
-    title="SAM 2 Everything Segmentation",
-    description="Upload an image to segment all objects using SAM 2."
-)
+with gr.Blocks() as demo:
+    gr.Markdown("# CLIP and FastSAM Demo")
+    
+    with gr.Tab("CLIP Zero-Shot Classification"):
+        with gr.Row():
+            image_input = gr.Image(type="pil", label="Input Image")
+            text_input = gr.Textbox(label="What do you want to check in the image?", placeholder="Type here...")
+        output_text = gr.Textbox(label="Result")
+        classify_btn = gr.Button("Classify")
+        classify_btn.click(fn=process_image_clip, inputs=[image_input, text_input], outputs=output_text)
+    
+    with gr.Tab("FastSAM Segmentation"):
+        with gr.Row():
+            image_input_sam = gr.Image(type="pil", label="Input Image")
+            image_output = gr.Image(type="pil", label="Segmentation Result")
+        segment_btn = gr.Button("Segment")
+        segment_btn.click(fn=process_image_fastsam, inputs=[image_input_sam], outputs=image_output)
 
-# Launch the interface
-iface.launch()
+if __name__ == "__main__":
+    demo.launch()