First attempt at serving model as a space

app.py

@@ -1,4 +1,7 @@
 import gradio as gr
+import zombie
+from huggingface_hub import hf_hub_download
+import onnxruntime as ort
 # import torch
 # from your_pix2pixhd_code import YourPix2PixHDModel, load_image, tensor2im # Adapt these imports
 
@@ -9,20 +12,31 @@ import gradio as gr
 # model.load_state_dict(torch.load('models/your_pix2pixhd_model.pth'))
 # model.eval()
 
+model_path = hf_hub_download(repo_id="jbrownkramer/makemeazombie", filename="smaller512x512_32bit.onnx")
+ort_session = ort.InferenceSession(model_path, providers=['CUDAExecutionProvider'])
+
 # --- 2. Define the prediction function ---
-def predict(input_image):
-    return input_image[..., ::-1]
-    # # Pre-process the input image
-    # processed_image = load_image(input_image)
+# def predict(input_image):
+#     return input_image[..., ::-1]
+#     # # Pre-process the input image
+#     # processed_image = load_image(input_image)
 
-    # # Run inference
-    # with torch.no_grad():
-    #     generated_image_tensor = model(processed_image)
+#     # # Run inference
+#     # with torch.no_grad():
+#     #     generated_image_tensor = model(processed_image)
 
-    # # Post-process the output tensor to an image
-    # output_image = tensor2im(generated_image_tensor)
+#     # # Post-process the output tensor to an image
+#     # output_image = tensor2im(generated_image_tensor)
 
-    # return output_image
+#     # return output_image
+
+def predict(input_image):
+    zombie_image = zombie.transition_onnx(input_image,ort_session)
+    
+    if zombie_image is None:
+        return "No face found"
+    
+    return zombie_image
 
 # --- 3. Create the Gradio Interface ---
 title = "pix2pixHD Image-to-Image Translation"
@@ -31,8 +45,8 @@ article = "<p style='text-align: center'>Model based on the <a href='https://git
 
 gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="numpy", label="Input Image"),
-    outputs=gr.Image(type="numpy", label="Output Image"),
+    inputs=gr.Image(type="pil", label="Input Image"),
+    outputs=gr.Image(type="pil", label="Output Image"),
     title=title,
     description=description,
     article=article,

requirements.txt

@@ -1 +1,3 @@
-gradio
+gradio
+onnxruntime-gpu
+opencv-python

zombie.py

@@ -0,0 +1,163 @@
+# import sys
+# sys.path.append("pix2pixHD")
+
+from PIL import Image
+import numpy as np
+# import pickle
+# import align
+# import time
+import cv2
+
+# sys.path.append(r"..\face-parsing.PyTorch")
+# import inferface
+
+# def get_model():
+#     with open("opt.pkl","rb") as f:
+#         opt = pickle.load(f)
+    
+#     return create_model(opt)
+
+def normalized_tensor(pil):
+    transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5),(0.5, 0.5, 0.5))])
+    return transform(pil)
+
+def tensor2im(tensor):
+    return Image.fromarray((255 * ((tensor.cpu().numpy() * .5) + .5)).astype(np.uint8).transpose((1,2,0)))
+
+def normalized_array(image):
+    np_array = np.array(image)
+    return np.expand_dims(((np_array / 255 - .5)/.5).astype(np.float32).transpose((2,0,1)),0)
+
+def array2im(np_array):
+    return Image.fromarray((255 * ((np_array * .5) + .5)).astype(np.uint8).transpose((1,2,0)))
+
+def square_human_2_zombie(image,model):
+    tensor = normalized_tensor(image)
+    generated = model(tensor.unsqueeze(0))
+    return tensor2im(generated.data[0])
+
+def square_human_2_zombie_onnx(image,model):
+    np_array = normalized_array(image)
+    ort_inputs = {"input": np_array}
+    ort_outs = model.run(None, ort_inputs)
+    return array2im(ort_outs[0][0])
+
+def human_2_zombie_onnx(rgb_image,model,output_size=512):
+    square,_ = align.align(rgb_image,enable_padding=False,output_size=output_size)
+    if square is None:
+        return None
+    return square_human_2_zombie_onnx(square,model)
+
+def human_2_zombie(rgb_image,model,output_size=512):
+    square,_ = align.align(rgb_image,enable_padding=False,output_size=output_size)
+    if square is None:
+        return None
+    return square_human_2_zombie(square,model)
+
+def human_path_2_zombie(path):
+    rgb_image = ImageOpen(path).convert("RGB")
+    return human_2_zombie(rgb_image)
+
+def transition(path,model):
+    rgb_image = ImageOpen(path).convert("RGB")
+    square,_ = align.align(rgb_image,enable_padding=False)
+    if square is None:
+        return None
+    zombie = square_human_2_zombie(square,model)
+    return get_concat_h(square,zombie)
+
+def transition_onnx(rgb_image,model):
+    #square,_ = align.align(rgb_image,enable_padding=False)
+    # if square is None:
+    #     return None
+
+    #Take the largest square in the upper left corner
+    #And resize it to 512x512
+    side = min(rgb_image.width,rgb_image.height)
+    square = rgb_image.crop((0,0,side,side))
+    square = square.resize((512,512))
+    
+    zombie = square_human_2_zombie_onnx(square,model)
+    
+    return get_concat_h(square,zombie)
+
+def get_concat_h(im1,im2):
+    dst = Image.new('RGB', (im1.width + im2.width, im1.height))
+    dst.paste(im1, (0, 0))
+    dst.paste(im2, (im1.width, 0))
+    return dst
+
+def get_concat_multiple_h(*ims):
+    to_return = ims[0]
+    for i in range(1,len(ims)):
+        to_return = get_concat_h(to_return,ims[i])
+        
+    return to_return
+
+def ImageOpen(filepath):
+    try:
+        image=Image.open(filepath)
+
+        for orientation in ExifTags.TAGS.keys():
+            if ExifTags.TAGS[orientation]=='Orientation':
+                break
+
+        exif=dict(image._getexif().items())
+
+        if exif[orientation] == 3:
+            image=image.rotate(180, expand=True)
+        elif exif[orientation] == 6:
+            image=image.rotate(270, expand=True)
+        elif exif[orientation] == 8:
+            image=image.rotate(90, expand=True)
+
+        image = image.convert('RGB')
+
+        return image
+    except (AttributeError, KeyError, IndexError):
+        # cases: image don't have getexif
+        return Image.open(filepath)
+
+def do_face(facenet, aligned, box, im_array, ort_session):
+        z = square_human_2_zombie_onnx(aligned,ort_session)
+        t1 = time.time()
+
+        t = cv2.getAffineTransform(np.array([[0,0],[0,511],[511,511]],dtype="float32"),box[:3,:].astype("float32"))
+
+        h,w,c = im_array.shape
+        face_in_place = cv2.warpAffine(np.array(z),t,(w,h))
+        
+        face_classes = inferface.run_net(facenet,aligned)
+        face_mask = np.logical_and(face_classes < 14, face_classes > 0)
+        
+        # face_classes_z = inferface.run_net(facenet,z)
+        # face_mask_z = np.logical_and(face_classes_z < 14, face_classes > 0)
+        
+        # face_mask = np.logical_or(face_mask,face_mask_z)
+
+        imagemask = cv2.warpAffine(face_mask.astype("uint8"),t,(w,h))
+
+        imagemask = imagemask.astype("uint8")
+
+        cv2.copyTo(face_in_place,imagemask,im_array)
+        
+def make_faces_zombie(path, facenet, ort_session):
+    im = Image.open(path)
+    im = im.convert(mode="RGB")
+    im_array = np.array(im)
+
+    return make_faces_zombie_from_array(im_array, facenet, ort_session)
+
+def make_faces_zombie_from_array(im_array_rgb, facenet, ort_session):
+    im_array_rgb = np.copy(im_array_rgb)
+
+    t0 = time.time()
+    faces = align.aligns(Image.fromarray(im_array_rgb),enable_padding=True,output_size=512)
+    print("Find faces",time.time() - t0)
+    if faces is None:
+        faces = []
+    
+    for aligned,box in faces:
+        do_face(facenet,aligned,box,im_array_rgb,ort_session)
+
+    return Image.fromarray(im_array_rgb)