rayuga2503 commited on
Commit
93c21cd
Β·
verified Β·
1 Parent(s): f3bf017

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +65 -0
app.py ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ import numpy as np
4
+ import cv2
5
+ from PIL import Image
6
+ from transformers import BlipProcessor, BlipForConditionalGeneration, SamModel, SamProcessor
7
+ import time
8
+
9
+ # Set device to GPU if available
10
+ device = "cuda" if torch.cuda.is_available() else "cpu"
11
+
12
+ # Load Florence BLIP model (Public Model - No Authentication Required)
13
+ processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
14
+ model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
15
+
16
+ # Load SAM model (Public Model - No Authentication Required)
17
+ sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
18
+ sam_model = SamModel.from_pretrained("facebook/sam-vit-base").to(device)
19
+
20
+ def process_image(image):
21
+ start_time = time.time()
22
+
23
+ # Convert and resize image
24
+ pil_image = Image.fromarray(image).resize((512, 512)) # Resize to optimize processing
25
+ print("βœ… Image loaded and resized.")
26
+
27
+ # Generate caption using Florence
28
+ try:
29
+ inputs = processor(pil_image, return_tensors="pt").to(device)
30
+ with torch.no_grad():
31
+ out = model.generate(**inputs)
32
+ description = processor.decode(out[0], skip_special_tokens=True)
33
+ print(f"πŸ“ Florence Captioning done in {time.time() - start_time:.2f} sec")
34
+ except Exception as e:
35
+ print(f"❌ Error in Florence: {e}")
36
+ return "Failed to generate description.", image
37
+
38
+ # Process Image for SAM
39
+ try:
40
+ encoding = sam_processor(images=pil_image, return_tensors="pt").to(device)
41
+ with torch.no_grad():
42
+ outputs = sam_model(**encoding)
43
+
44
+ # Extract segmentation mask
45
+ mask = outputs.pred_masks[0, 0].cpu().numpy()
46
+ mask_overlay = image.copy()
47
+ mask_overlay[mask > 0.5] = [0, 255, 0] # Green overlay for segmentation
48
+ print(f"🎨 SAM Segmentation done in {time.time() - start_time:.2f} sec")
49
+ except Exception as e:
50
+ print(f"❌ Error in SAM: {e}")
51
+ return description, image
52
+
53
+ return description, mask_overlay
54
+
55
+ # Gradio Interface
56
+ demo = gr.Interface(
57
+ fn=process_image,
58
+ inputs=gr.Image(type="numpy"),
59
+ outputs=[gr.Textbox(label="Image Description"), gr.Image(label="Segmented Image")],
60
+ title="Florence + SAM Image Processing",
61
+ description="Upload an image to get its description using Florence and segmentation using SAM (loaded from Hugging Face)."
62
+ )
63
+
64
+ if __name__ == "__main__":
65
+ demo.launch(server_name="0.0.0.0", server_port=7860)