2024.07.04 update app.py

.ipynb_checkpoints/app-checkpoint.py

@@ -0,0 +1,82 @@
+#  객체검출 -> 삭제 체크박스 적용본
+import torch
+from PIL import Image, ImageDraw
+from transformers import DetrImageProcessor, DetrForObjectDetection
+from diffusers import StableDiffusionInpaintPipeline
+import gradio as gr
+
+# 모델 로드
+processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
+model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16)
+pipe = pipe.to("cpu")
+
+def detect_objects(image):
+    # 객체 검출
+    inputs = processor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+
+    # 결과 후처리```````
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)[0]
+
+    # 검출된 객체 정보 추출
+    detected_objects = []
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        if score > 0.9:
+            box = [round(i) for i in box.tolist()]
+            detected_objects.append({"label": model.config.id2label[label.item()], "box": box})
+    
+    return detected_objects
+
+def display_detected_objects(image):
+    detected_objects = detect_objects(image)
+    labeled_image = image.copy()
+    draw = ImageDraw.Draw(labeled_image)
+    object_labels = []
+    for obj in detected_objects:
+        box = obj["box"]
+        label = obj["label"]
+        draw.rectangle(box, outline="red", width=3)
+        draw.text((box[0], box[1]), label, fill="red")
+        object_labels.append(f"{label} at {box}")
+    return labeled_image, gr.update(choices=object_labels)
+
+def inpaint_image(image, selected_objects):
+    detected_objects = detect_objects(image)
+
+    # 마스크 생성
+    mask = Image.new("L", image.size, 0)
+    draw = ImageDraw.Draw(mask)
+    for obj in detected_objects:
+        object_label = f"{obj['label']} at {obj['box']}"
+        if object_label in selected_objects:
+            box = obj["box"]
+            draw.rectangle(box, fill=255)
+    
+    # Inpainting 수행
+    image = image.convert("RGB")
+    mask = mask.convert("RGB")
+    output = pipe(prompt="a modern interior", image=image, mask_image=mask).images[0]
+    # output = pipe(prompt="remove", image=image, mask_image=mask).images[0]
+
+    
+    return output
+
+# Gradio 인터페이스 설정
+with gr.Blocks() as interface:
+    with gr.Row():
+        image_input = gr.Image(type="pil", label="Input Image")
+        objects_list = gr.CheckboxGroup(label="Detected Objects")
+    
+    labeled_image_output = gr.Image(label="Labeled Image")
+    final_output = gr.Image(label="Output Image")
+    
+    detect_button = gr.Button("Detect Objects")
+    inpaint_button = gr.Button("Remove Selected Objects")
+    
+    detect_button.click(fn=display_detected_objects, inputs=image_input, outputs=[labeled_image_output, objects_list])
+    inpaint_button.click(fn=inpaint_image, inputs=[image_input, objects_list], outputs=final_output)
+
+# Gradio 인터페이스 실행
+interface.launch()

app.py

@@ -8,7 +8,7 @@ import gradio as gr
 # 모델 로드
 processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
 model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
-pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting", torch_dtype=torch.float32)
+pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16)
 pipe = pipe.to("cpu")
 
 def detect_objects(image):