| import io | |
| import json | |
| import os | |
| import tempfile | |
| import zipfile | |
| import datetime | |
| import requests | |
| import random | |
| import gradio as gr | |
| from PIL import Image, ImageDraw, ImageFont | |
| API_URL = os.getenv("API_URL") | |
| API_KEY = os.getenv("API_KEY") | |
| MODEL = os.getenv("MODEL") | |
| ALLOWED_EXTENSIONS = ( | |
| ".png", ".jpg", ".jpeg", ".bmp", ".gif", ".webp", | |
| ".tif", ".tiff", ".avif", ".ico", ".ppm", ".pgm", ".pbm", ".eps", ".icns" | |
| ) | |
| def get_color_for_label(label, color_mapping): | |
| if label in color_mapping: | |
| return color_mapping[label] | |
| color = (random.randint(0,255), random.randint(0,255), random.randint(0,255)) | |
| color_mapping[label] = color | |
| return color | |
| def process_single_image(image, prompt, confidence, file_format): | |
| if "," in prompt: | |
| prompt_list = [p.strip() for p in prompt.split(",")] | |
| else: | |
| prompt_list = [prompt.strip()] | |
| prompt_list_lower = [p.lower() for p in prompt_list] | |
| buffered = io.BytesIO() | |
| image.save(buffered, format="PNG") | |
| buffered.seek(0) | |
| files = {"image": buffered} | |
| data = { | |
| "prompts": prompt_list, | |
| "model": MODEL, | |
| "confidence": confidence | |
| } | |
| headers = {"Authorization": "Basic " + API_KEY} | |
| response = requests.post(API_URL, files=files, data=data, headers=headers) | |
| response_data = response.json() | |
| print("Response data:", response_data) | |
| detections_list = response_data.get("data", []) | |
| if detections_list and isinstance(detections_list, list): | |
| detections = detections_list[0] | |
| else: | |
| detections = [] | |
| annotated_image = image.copy() | |
| draw = ImageDraw.Draw(annotated_image) | |
| try: | |
| font = ImageFont.truetype("arial.ttf", size=15) | |
| except IOError: | |
| font = ImageFont.load_default() | |
| color_mapping = {} | |
| for det in detections: | |
| box = det["bounding_box"] | |
| score = det["score"] | |
| detection_label = det.get("label", prompt).strip() | |
| box_color = get_color_for_label(detection_label, color_mapping) | |
| draw.rectangle(box, outline=box_color, width=2) | |
| text = f"{detection_label} {score:.2f}" | |
| bbox = draw.textbbox((0, 0), text, font=font) | |
| text_width = bbox[2] - bbox[0] | |
| text_height = bbox[3] - bbox[1] | |
| text_origin = (box[0], box[1] - text_height) | |
| draw.rectangle([text_origin, (box[0] + text_width, box[1])], fill=box_color) | |
| draw.text(text_origin, text, fill="white", font=font) | |
| img_width, img_height = image.size | |
| if file_format.lower() == "json": | |
| detection_info = {"detections": detections} | |
| ext = ".json" | |
| content = json.dumps(detection_info, indent=2) | |
| elif file_format.lower() == "txt": | |
| lines = [] | |
| for det in detections: | |
| detection_label = det.get("label", prompt).strip().lower() | |
| if detection_label not in prompt_list_lower: | |
| continue | |
| class_id = prompt_list_lower.index(detection_label) | |
| x1, y1, x2, y2 = det["bounding_box"] | |
| x_center = ((x1 + x2) / 2) / img_width | |
| y_center = ((y1 + y2) / 2) / img_height | |
| bbox_width = (x2 - x1) / img_width | |
| bbox_height = (y2 - y1) / img_height | |
| line = f"{class_id} {x_center:.6f} {y_center:.6f} {bbox_width:.6f} {bbox_height:.6f}" | |
| lines.append(line) | |
| content = "\n".join(lines) | |
| ext = ".txt" | |
| else: | |
| detection_info = {"detections": detections} | |
| ext = ".json" | |
| content = json.dumps(detection_info, indent=2) | |
| temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=ext, mode="w", encoding="utf-8") | |
| temp_file.write(content) | |
| temp_file.close() | |
| return annotated_image, temp_file.name | |
| def auto_annotate_batch(input_files, prompt, confidence, file_format): | |
| images_and_names = [] | |
| if len(input_files) == 1 and str(input_files[0]).lower().endswith(".zip"): | |
| zip_path = input_files[0] | |
| extract_dir = tempfile.mkdtemp() | |
| with zipfile.ZipFile(zip_path, 'r') as zip_ref: | |
| zip_ref.extractall(extract_dir) | |
| for root, _, files in os.walk(extract_dir): | |
| for file in files: | |
| if file.lower().endswith(ALLOWED_EXTENSIONS): | |
| img_path = os.path.join(root, file) | |
| try: | |
| img = Image.open(img_path).convert("RGB") | |
| images_and_names.append((img, file)) | |
| except Exception as e: | |
| print(f"Gagal membuka {img_path}: {e}") | |
| else: | |
| for file_path in input_files: | |
| if file_path.lower().endswith(ALLOWED_EXTENSIONS): | |
| try: | |
| img = Image.open(file_path).convert("RGB") | |
| original_filename = os.path.basename(file_path) | |
| images_and_names.append((img, original_filename)) | |
| except Exception as e: | |
| print(f"Gagal membuka {file_path}: {e}") | |
| annotated_images = [] | |
| detection_file_paths = [] | |
| for img, original_filename in images_and_names: | |
| ann_img, temp_det_file = process_single_image(img, prompt, confidence, file_format) | |
| annotated_images.append(ann_img) | |
| base_name, _ = os.path.splitext(original_filename) | |
| new_name = base_name + (".json" if file_format.lower() == "json" else ".txt") | |
| new_path = os.path.join(os.path.dirname(temp_det_file), new_name) | |
| os.rename(temp_det_file, new_path) | |
| detection_file_paths.append(new_path) | |
| timestamp = datetime.datetime.now().strftime("%d-%m-%Y_%H-%M-%S") | |
| zip_filename = f"Data_Annotate_{timestamp}.zip" | |
| zip_out_path = os.path.join(tempfile.gettempdir(), zip_filename) | |
| with zipfile.ZipFile(zip_out_path, 'w') as zipf: | |
| for file_path in detection_file_paths: | |
| zipf.write(file_path, os.path.basename(file_path)) | |
| return annotated_images, zip_out_path, detection_file_paths | |
| iface = gr.Interface( | |
| fn=auto_annotate_batch, | |
| inputs=[ | |
| gr.File(file_count="multiple", type="filepath", label="Upload Image (multiple files or ZIP)"), | |
| gr.Textbox(lines=1, placeholder="Enter object prompt (separate with comma if more than one)", label="Prompt"), | |
| gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.25, label="Confidence Threshold"), | |
| gr.Radio(choices=["json", "txt"], value="json", label="Format Annotation File") | |
| ], | |
| outputs=[ | |
| gr.Gallery(label="Annotated Images"), | |
| gr.File(label="Download All Annotations (ZIP)"), | |
| gr.File(label="Download Individual Annotation Files") | |
| ], | |
| title="Auto Annotate") | |
| if __name__ == "__main__": | |
| iface.launch(debug=True) |