Update app.py

app.py

@@ -1,8 +1,12 @@
 from run import process
 import time
+import subprocess
 import os
+import argparse
 import cv2
+import sys
 from PIL import Image
+import torch
 import gradio as gr
 
 TESTdevice = "cpu"
@@ -36,12 +40,27 @@ def inference(img):
     return output
 
 def load_image_from_file(file_path, new_height=None):
+    """
+    Load an image from a file and optionally resize it while maintaining the aspect ratio.
+
+    Args:
+        file_path (str): The path to the image file.
+        new_height (int, optional): The new height for the image. If None, the image is not resized.
+
+    Returns:
+        Image: The loaded (and optionally resized) image.
+    """
     try:
         img = Image.open(file_path)
-        if new_height is not None:
+        
+        if (new_height is not None):
+            # Calculate new width to maintain aspect ratio
             aspect_ratio = img.width / img.height
             new_width = int(new_height * aspect_ratio)
+            
+            # Resize the image
             img = img.resize((new_width, new_height), Image.LANCZOS)
+        
         return img
     except FileNotFoundError:
         print(f"File not found: {file_path}")
@@ -65,12 +84,104 @@ examples = [
     [load_image_from_file('example8.webp')],
 ]
 
-with gr.Blocks(theme="Nymbo/Alyx_Theme") as demo:
-    width = 240
-    height = 340
+js='''
+<script>
+window.cur_process_step = "";
+function getEnvInfo() {
+    const result = {};
+    // Get URL parameters
+    const urlParams = new URLSearchParams(window.location.search);
+    for (const [key, value] of urlParams) {
+        result[key] = value;
+    }
+    // Get current domain and convert to lowercase
+    result["__domain"] = window.location.hostname.toLowerCase();
+    // Get iframe parent domain, if any, and convert to lowercase
+    try {
+        if (window.self !== window.top) {
+            result["__iframe_domain"] = document.referrer 
+                ? new URL(document.referrer).hostname.toLowerCase()
+                : "unable to get iframe parent domain";
+        }else{
+            result["__iframe_domain"] = "";
+        }
+    } catch (e) {
+        result["__iframe_domain"] = "unable to access iframe parent domain";
+    }
+    return result;
+}
+function isValidEnv(){
+    envInfo = getEnvInfo();
+    return envInfo["e"] == "1" || 
+        envInfo["__domain"].indexOf("nsfwais.io") != -1 || 
+        envInfo["__iframe_domain"].indexOf("nsfwais.io") != -1 ||
+        envInfo["__domain"].indexOf("127.0.0.1") != -1 || 
+        envInfo["__iframe_domain"].indexOf("127.0.0.1") != -1;
+}
+window.postMessageToParent = function(img, event, source, value) {
+    // Construct the message object with the provided parameters
+    console.log("post start",event, source, value);
+    const message = {
+        event: event,
+        source: source,
+        value: value
+    };
+    
+    // Post the message to the parent window
+    window.parent.postMessage(message, '*');
+    console.log("post finish");
+    window.cur_process_step = "process";
+    return img;
+}
+function uploadImage(image, event, source, value) {
+    // Ensure we're in an iframe
+    if (window.cur_process_step != "process"){
+        return;
+    }
+    window.cur_process_step = "";
+    console.log("uploadImage", image ? image.url : null, event, source, value);
+    // Get the first image from the gallery (assuming it's an array)
+    let imageUrl = image ? image.url : null;
+    if (window.self !== window.top) {
+        // Post the message to the parent window
+        // Prepare the data to send
+        let data = {
+            event: event,
+            source: source,
+            value: imageUrl
+        };
+        window.parent.postMessage(data, '*');
+    } else if (isValidEnv()){
+        try{
+            sendCustomEventToDataLayer({},event,source,{"image":imageUrl})
+        } catch (error) {
+            console.error("Error in sendCustomEventToDataLayer:", error);
+        }
+    }else{
+        console.log("Not in an iframe, can't post to parent");
+    }
+    return;
+}
+window.onDemoLoad = function(x){
+    let envInfo = getEnvInfo();
+    console.log(envInfo);
+    if (isValidEnv()){
+        var element = document.getElementById("pitch_desc_html_code");
+        if (element) {
+            element.parentNode.removeChild(element);
+        }
+    }
+    return "";
+}
+</script>
+'''
+
+with gr.Blocks(head=js, theme="Nymbo/Alyx_Theme") as demo:
+    width=240
+    height=340
 
     with gr.Row(equal_height=False):
-        with gr.Column(min_width=240):
+        with gr.Column(min_width=240):  # Adjust scale for proper sizing
             image_input = gr.Image(type="numpy", label="", height=height)
             gr.Examples(examples=examples, inputs=image_input, examples_per_page=10, elem_id="example_img")
             process_button = gr.Button("Nude!", size="sm")
@@ -79,8 +190,8 @@ with gr.Blocks(theme="Nymbo/Alyx_Theme") as demo:
         processed_img = inference(img)
         return processed_img
 
-    image_input.change(fn=lambda x: x, inputs=[image_input], outputs=[gr.State([])])
-    process_button.click(update_status, inputs=image_input, outputs=image_input)
-
+    image_input.change(fn=lambda x: x, inputs=[image_input], outputs=[gr.State([])], js='''(img) => window.uploadImage(img, "process_finished", "demo_hf_deepnude_gan_card", "")''')
+    process_button.click(update_status, inputs=image_input, outputs=image_input, js='''(i) => window.postMessageToParent(i, "process_started", "demo_hf_deepnude_gan_card", "click_nude")''')
+    demo.load(fn=lambda x: x, inputs=[gr.State([])], outputs=[gr.State([])], js='''(x) => window.onDemoLoad(x)''')
 demo.queue(max_size=10)
 demo.launch()