import matplotlib.pyplot as plt import numpy as np from collections import defaultdict from ultralytics import YOLO import gradio as gr from PIL import Image import io import base64 # import matplotlib.pyplot as plt from transformers import pipeline # import numpy as np # from collections import defaultdict # from ultralytics import YOLO # import gradio as gr # from PIL import Image MODEL_PATH = 'best-1206.pt' model = YOLO(MODEL_PATH) checkpoint = "openai/clip-vit-base-patch32" classifier = pipeline(model=checkpoint, task="zero-shot-image-classification") def shot(image, labels_text): image = Image.fromarray(np.uint8(image)).convert('RGB') if labels_text: labels = labels_text.split(";") else: labels = ['PET bottle', 'Metal can', 'Plastic bag', 'Cardboard box'] results = classifier(image, candidate_labels=labels) return {result["label"]: result["score"] for result in results} def count_and_display_pie_chart(image, labels_text): image = Image.fromarray(np.uint8(image)).convert('RGB') if labels_text: labels = labels_text.split(";") else: labels = ['PET bottle', 'Metal can', 'Plastic bag', 'Cardboard box'] results = model(image) boxes = results[0].boxes.xyxy.tolist() classes = results[0].boxes.cls.tolist() names = results[0].names confidences = results[0].boxes.conf.tolist() label_count = defaultdict(int) for i, (box, cls, conf) in enumerate(zip(boxes, classes, confidences)): x1, y1, x2, y2 = box confidence = conf detected_class = cls name = names[int(cls)] # Crop the bounding box from the image #cropped_image = image[int(y1):int(y2), int(x1):int(x2)] cropped_image = image.crop((int(x1), int(y1), int(x2), int(y2))) # Pass the cropped image through the 'shot' function #labels_text = 'PET bottle, Milk container, Soda can, Plastic cover' # Assuming 'name' contains the label text for the detected object result = shot(cropped_image, labels_text) # Find the label with the highest score max_label = max(result, key=result.get) max_score = result[max_label] if max_score > 0.5: # Increment the count for the detected label label_count[max_label] += 1 labels = list(label_count.keys()) sizes = list(label_count.values()) fig, ax = plt.subplots() ax.pie(sizes, labels=labels, autopct='%1.1f%%', startangle=140) ax.axis('equal') ax.set_title('Detected Objects Pie Chart') buffer = io.BytesIO() fig.savefig(buffer, format='png') buffer.seek(0) img_str = base64.b64encode(buffer.read()).decode('utf-8') img_html = ''.format(img_str) return img_html demo = gr.Interface( count_and_display_pie_chart, inputs=[ gr.Image(type="numpy", label="Upload an image"), gr.Textbox( label="Labels", info="Separated by a semicolon (;)", lines=4, value="""PET bottle; Metal can; Plastic bag; Cardboard box""", ) ], outputs=[gr.HTML(label="Pie Chart Image")], examples=[['can.jpg', None]], description="Upload an image to detect objects and display a pie chart of their counts.", title="Object Detection and Pie Chart Demo" ) demo.launch() # import matplotlib.pyplot as plt # from transformers import pipeline # import numpy as np # from collections import defaultdict # from ultralytics import YOLO # import gradio as gr # from PIL import Image # MODEL_PATH='best-1206.pt' # model = YOLO(MODEL_PATH) # pipe = pipeline("zero-shot-image-classification", model="openai/clip-vit-base-patch32") # checkpoint = "openai/clip-vit-base-patch32" # classifier = pipeline(model=checkpoint, # task="zero-shot-image-classification") # def shot(image, labels_text): # image = Image.fromarray(np.uint8(image)).convert('RGB') # if labels_text: # labels = labels_text.split(";") # else: # labels = ['PET bottle', # 'Metal can', # 'Plastic bag', # 'Cardboard box'] # results = classifier(image, # candidate_labels=labels) # return {result["label"]: result["score"] for result in results} # # # Create a defaultdict to count occurrences of each label # # label_count = defaultdict(int) # # # Iterate through the results # # for i, (box, cls, conf) in enumerate(zip(boxes, classes, confidences)): # # x1, y1, x2, y2 = box # # confidence = conf # # detected_class = cls # # name = names[int(cls)] # # cropped_image = image.crop((int(x1), int(y1), int(x2), int(y2))) # # # For simplicity, assume 'labels_text' remains unchanged from the function arguments # # result = shot(np.array(cropped_image), labels_text) # # max_label = max(result, key=result.get) # # max_score = result[max_label] # # if max_score > 0.5: # # label_count[max_label] += 1 # # # Generate a pie chart based on the label counts # # labels = list(label_count.keys()) # # sizes = list(label_count.values()) # # fig, ax = plt.subplots() # # ax.pie(sizes, labels=labels, autopct='%1.1f%%', startangle=140) # # ax.axis('equal') # Equal aspect ratio ensures that pie is drawn as a circle. # # ax.set_title('Detected Objects Pie Chart') # # return fig # Return the Matplotlib figure # # # # Create a defaultdict to count occurrences of each label # # # label_count = defaultdict(int) # # # # Iterate through the results # # # for i, (box, cls, conf) in enumerate(zip(boxes, classes, confidences)): # # # x1, y1, x2, y2 = box # # # confidence = conf # # # detected_class = cls # # # name = names[int(cls)] # # # # Crop the bounding box from the image # # # #cropped_image = image[int(y1):int(y2), int(x1):int(x2)] # # # cropped_image = image.crop((int(x1), int(y1), int(x2), int(y2))) # # # # Pass the cropped image through the 'shot' function # # # labels_text = 'PET bottle, Milk container, Soda can, Plastic cover' # Assuming 'name' contains the label text for the detected object # # # result = shot(cropped_image, labels_text) # # # # Find the label with the highest score # # # max_label = max(result, key=result.get) # # # max_score = result[max_label] # # # if max_score > 0.5: # # # # Increment the count for the detected label # # # label_count[max_label] += 1 # # # return label_count # # # demo = gr.Interface(count, # # # [gr.Image(type="pil"), # # # gr.Textbox( # # # label="Labels", # # # info="Separated by a semicolon (;)", # # # lines=4, # # # value="""PET bottle; # # # Metal can; # # # Plastic bag; # # # Cardboard box""", # # # )], # # # outputs="label", # # # examples=[['can.jpg',None]], # # # description="Upload an image", # # # title="Recycling AI Data Demo") # # # demo.launch() # # demo = gr.Interface( # # count_and_display_pie_chart, # # inputs=[ # # gr.Image(type="numpy", label="Upload an image"), # # gr.Textbox( # # label="Labels", # # info="Separated by a semicolon (;)", # # lines=4, # # value="""PET bottle; # # Metal can; # # Plastic bag; # # Cardboard box""", # # ) # # ], # # outputs=gr.outputs.Image(type="plot", label="Pie Chart"), # # examples=[['test.png',None]], # # description="Upload an image to detect objects and display a pie chart of their counts.", # # title="Object Detection and Pie Chart Demo" # # ) # # demo.launch()