SahilCarterr commited on
Commit
122ad84
·
verified ·
1 Parent(s): eaccd6c

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +94 -0
app.py ADDED
@@ -0,0 +1,94 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ! git clone https://github.com/mikonvergence/ControlNetInpaint
2
+ import sys
3
+ sys.path.append('./ControlNetInpaint/')
4
+ from ultralytics import YOLO
5
+ from PIL import Image
6
+ import gradio as gr
7
+ import numpy as np
8
+ from diffusers import StableDiffusionInpaintPipeline, ControlNetModel, UniPCMultistepScheduler
9
+ from diffusers.utils import load_image
10
+ pipe_sd = StableDiffusionInpaintPipeline.from_pretrained(
11
+ "runwayml/stable-diffusion-inpainting",
12
+ revision="fp16",
13
+ torch_dtype=torch.float16,
14
+ )
15
+ # speed up diffusion process with faster scheduler and memory optimization
16
+ pipe_sd.scheduler = UniPCMultistepScheduler.from_config(pipe_sd.scheduler.config)
17
+ #pipe_sd.to('cpu')
18
+ # load control net and stable diffusion v1-5
19
+ from diffusers import StableDiffusionControlNetInpaintPipeline
20
+ controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
21
+ pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
22
+ "runwayml/stable-diffusion-inpainting", controlnet=controlnet, torch_dtype=torch.float16
23
+ )
24
+
25
+ # speed up diffusion process with faster scheduler and memory optimization
26
+ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
27
+ # remove following line if xformers is not installed
28
+ text_prompt="Transform this image into a work of art by changing its style and color palette. Apply a distinct artistic style, such as impressionism, cubism, or surrealism, to give the image a unique and visually striking appearance. Experiment with brush strokes, textures, and effects to achieve the desired artistic effect while maintaining the essence of the original scene. Additionally, adjust the color palette to evoke a specific mood or theme. For example, infuse warm, earthy tones for a rustic and cozy feel, or opt for vibrant, psychedelic colors for a surreal and otherworldly atmosphere. The goal is to reimagine the image in a creative and expressive way, transforming it into a captivating visual masterpiece."
29
+ import os
30
+ os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
31
+ #pipe.to('cuda')
32
+ import cv2
33
+ from PIL import Image
34
+ import numpy as np
35
+ import torch
36
+ from matplotlib import pyplot as plt
37
+ def image_enhance(image, mask):
38
+ image = np.array(image)
39
+ mask_image = np.array(mask)
40
+ canny_image = cv2.Canny(image, 100, 200)
41
+ canny_image = canny_image[:, :, None]
42
+ canny_image = np.concatenate([canny_image, canny_image, canny_image], axis=2)
43
+
44
+ image=Image.fromarray(image)
45
+ mask_image=Image.fromarray(mask_image)
46
+ canny_image = Image.fromarray(canny_image)
47
+
48
+ # generate image
49
+ generator = torch.manual_seed(42)
50
+ new_image = pipe(
51
+ text_prompt,
52
+ num_inference_steps=20,
53
+ generator=generator,
54
+ image=image,
55
+ control_image=canny_image,
56
+ controlnet_conditioning_scale = 0.5,
57
+ mask_image=mask_image
58
+ ).images[0]
59
+
60
+ return new_image
61
+ def function(image):
62
+ original_image=image
63
+ model = YOLO('best_yolo_2.pt')
64
+ image = np.array(image)
65
+ results = model.predict(image)
66
+ img=Image.fromarray(results[0].plot())
67
+ num_result=len(results[0].boxes.cls)
68
+ names_output=[]
69
+ for i in range(num_result):
70
+ name=results[0].names[int(results[0].boxes.cls[i])]
71
+ names_output.append(name)
72
+ boxes = results[0].boxes
73
+ boxes=boxes.xyxy
74
+ #image=Image.fromarray(image)
75
+ image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
76
+ predictor = SamPredictor(sam)
77
+ predictor.set_image(image)
78
+ list_mask=[]
79
+ for box in boxes:
80
+ input_box = np.array(box.cpu())
81
+ masks, _, _ = predictor.predict(
82
+ point_coords=None,
83
+ point_labels=None,
84
+ box=input_box[None, :],
85
+ multimask_output=False,
86
+ )
87
+ for mask in masks:
88
+ list_mask.append(mask)
89
+ for i in range(len(list_mask)):
90
+ list_mask[i]=Image.fromarray(list_mask[i])
91
+ original_image=image_enhance(original_image,list_mask[i])
92
+ return original_image
93
+ demo=gr.Interface(fn=function,inputs="image",outputs=["image"])
94
+ demo.launch()