0llheaven commited on
Commit
a8172fd
·
verified ·
1 Parent(s): 4dc1d6d

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +53 -0
app.py ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoImageProcessor, AutoModelForObjectDetection
3
+ import torch
4
+ from PIL import Image, ImageDraw
5
+
6
+ # Load the model and processor
7
+ processor = AutoImageProcessor.from_pretrained("0llheaven/Conditional-detr-finetuned-tf")
8
+ model = AutoModelForObjectDetection.from_pretrained("0llheaven/Conditional-detr-finetuned-tf")
9
+
10
+ def detect_objects(image, score_threshold):
11
+ # Convert image to RGB if it's grayscale
12
+ if image.mode != "RGB":
13
+ image = image.convert("RGB")
14
+
15
+ # Prepare input for the model
16
+ inputs = processor(images=image, return_tensors="pt")
17
+ outputs = model(**inputs)
18
+
19
+ # Filter predictions based on the user-defined score threshold
20
+ target_sizes = torch.tensor([image.size[::-1]])
21
+ results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)
22
+
23
+ labels_output = []
24
+
25
+ # Draw bounding boxes around detected objects
26
+ draw = ImageDraw.Draw(image)
27
+ for result in results:
28
+ scores = result["scores"]
29
+ labels = result["labels"]
30
+ boxes = result["boxes"]
31
+
32
+ for score, label, box in zip(scores, labels, boxes):
33
+ if score >= score_threshold: # Only draw if score is above threshold
34
+ box = [round(i, 2) for i in box.tolist()]
35
+ label_name = "Pneumonia" if label.item() == 0 else "Other"
36
+ draw.rectangle(box, outline="red", width=3)
37
+ draw.text((box[0], box[1]), f"{label_name}: {round(score.item(), 3)}", fill="red")
38
+ labels_output.append(f"{label_name}: {round(score.item(), 3)}")
39
+
40
+ return image, "\n".join(labels_output)
41
+
42
+ # Create the Gradio interface
43
+ interface = gr.Interface(
44
+ fn=detect_objects,
45
+ inputs=[gr.Image(type="pil"), gr.Slider(0, 1, value=0.5, label="Score Threshold")], # Add slider for score threshold
46
+ # outputs=gr.Image(type="pil"), # Corrected output type
47
+ outputs=[gr.Image(type="pil"), gr.Textbox(label="Detected Objects")],
48
+ title="Object Detection with Transformers",
49
+ description="Upload an image to detect objects using a fine-tuned Conditional-DETR model."
50
+ )
51
+
52
+ # Launch the interface
53
+ interface.launch()