Update app.py

app.py

@@ -63,77 +63,77 @@ transform = T.Compose([
 
 def gen_sources(deepfake_img):
     #----------------DeepFake Face Segmentation-----------------
-    ##----------------------Initialize:Face Segmentation----------------------------------
     segmenter = FaceSegmenter(threshold=0.5)
-    # Convert PIL Image to BGR numpy array for segmentation
     img_np = np.array(deepfake_img.convert('RGB'))
     img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
-    # Segment the face
     segmented_np = segmenter.segment_face(img_bgr)
-    # Convert segmented numpy array (BGR) back to PIL Image
     deepfake_seg = Image.fromarray(cv2.cvtColor(segmented_np, cv2.COLOR_BGR2RGB))
-    #------------Initialize:Decoder-F------------------------
+
+    #------------Initialize Models------------------------
     checkpoint_path_f = "./models/model_vaq1_ff.pth"
-    # Load model checkpoints
     checkpoint_f = torch.load(checkpoint_path_f, map_location=device)
-    # Load the state dictionary into the models
     model_vaq_f.load_state_dict(checkpoint_f, strict=True)
     model_vaq_f.eval()
-    #------------Initialize:Decoder-G------------------------
+
     checkpoint_path_g = "./models/model_vaq2_gg.pth"
     checkpoint_g = torch.load(checkpoint_path_g, map_location=device)
-    # Load the state dictionary into the models
     model_vaq_g.load_state_dict(checkpoint_g, strict=True)
     model_vaq_g.eval()
-    ##------------------------Initialize Model-F-------------------------------------
+
     model_z1 = DeepfakeToSourceTransformer().to(device)
-    model_z1.load_state_dict(torch.load("./models/model_z1_ff.pth",map_location=device),strict=True)
+    model_z1.load_state_dict(torch.load("./models/model_z1_ff.pth", map_location=device), strict=True)
     model_z1.eval()
-    ##------------------------Initialize Model-G-------------------------------------
+
     model_z2 = DeepfakeToSourceTransformer().to(device)
-    model_z2.load_state_dict(torch.load("./models/model_z2_gg.pth",map_location=device),strict=True)
+    model_z2.load_state_dict(torch.load("./models/model_z2_gg.pth", map_location=device), strict=True)
     model_z2.eval()
-    ##--------------------Initialize:Evaluation---------------------------------------
+
     criterion = DF()
-    
-    ##----------------------Operation-------------------------------------------------
+
     with torch.no_grad():
-        # Load and preprocess input image
-        #img = Image.open(deepfake_img).convert('RGB')
-        df_img = transform(deepfake_img.convert('RGB')).unsqueeze(0).to(device)  # Shape: (1, 3, 256, 256)
+        df_img = transform(deepfake_img.convert('RGB')).unsqueeze(0).to(device)
         seg_img = transform(deepfake_seg).unsqueeze(0).to(device)
-        
-        # Calculate quantized_block for all images
-        z_df, _, _ = model_vaq_f.encode(df_img) 
-        z_seg, _, _ = model_vaq_g.encode(seg_img) 
-        rec_z_img1 = model_z1(z_df) 
-        rec_z_img2 = model_z2(z_seg) 
+
+        z_df, _, _ = model_vaq_f.encode(df_img)
+        z_seg, _, _ = model_vaq_g.encode(seg_img)
+        rec_z_img1 = model_z1(z_df)
+        rec_z_img2 = model_z2(z_seg)
         rec_img1 = model_vaq_f.decode(rec_z_img1).squeeze(0)
         rec_img2 = model_vaq_g.decode(rec_z_img2).squeeze(0)
         rec_img1_pil = T.ToPILImage()(rec_img1)
         rec_img2_pil = T.ToPILImage()(rec_img2)
 
-        # Save PIL images to in-memory buffers
-        buffer1 = BytesIO()
-        buffer2 = BytesIO()
-        rec_img1_pil.save(buffer1, format="PNG")
-        rec_img2_pil.save(buffer2, format="PNG")
+        # Save PIL images to temporary files
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp1, \
+             tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp2:
+            
+            rec_img1_pil.save(temp1, format="PNG")
+            rec_img2_pil.save(temp2, format="PNG")
+            
+            temp1_path = temp1.name
+            temp2_path = temp2.name
 
-        # Pass buffers to Gradio client
+        # Pass file paths to Gradio client
         result = client.predict(
-            target=file(buffer1),
-            source=file(buffer2), slider=100, adv_slider=100,
+            target=file(temp1_path),
+            source=file(temp2_path), slider=100, adv_slider=100,
             settings=["Adversarial Defense"], api_name="/run_inference"
         )
 
+        # Clean up temporary files
+        os.remove(temp1_path)
+        os.remove(temp2_path)
+
         # Load result and compute loss
-        dfimage_pil = Image.open(result)  # Open the resulting image
-        buffer3 = BytesIO()
-        dfimage_pil.save(buffer3, format="PNG")
-        rec_df = transform(Image.open(buffer3)).unsqueeze(0).to(device)
-        rec_loss,_ = criterion(df_img, rec_df)
+        dfimage_pil = Image.open(result)
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp3:
+            dfimage_pil.save(temp3, format="PNG")
+            rec_df = transform(Image.open(temp3.name)).unsqueeze(0).to(device)
+            os.remove(temp3.name)
+
+        rec_loss, _ = criterion(df_img, rec_df)
 
-        return (rec_img1_pil, rec_img2_pil, dfimage_pil, round(rec_loss.item(),3))
+        return (rec_img1_pil, rec_img2_pil, dfimage_pil, round(rec_loss.item(), 3))
 
 #________________________Create the Gradio interface_________________________________
 interface = gr.Interface(