diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..499ca79b7b466e5b556dca990f58f927acaa6b71 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +venv +__pycache__/ \ No newline at end of file diff --git a/app.py b/app.py index f337155d01cba37d90d3cfcacd8d70051f1a98f6..dbf066eb2a1427cfc160f6867cd5015a0c49cefa 100644 --- a/app.py +++ b/app.py @@ -1,5 +1,5 @@ import gradio as gr -import spaces +# import spaces import torch import torch.nn.functional as F from torch.utils.data import DataLoader @@ -14,17 +14,17 @@ from io import BytesIO # モデルとデータの読み込み def load_model(): - model_path = "checkpoints/ae_model_tf_2024-03-05_00-35-21.pth" - feature_dim = 32 + model_path = "checkpoints/autoencoder-epoch=49-train_loss=1.01.ckpt" + feature_dim = 64 model = AutoencoderModule(feature_dim=feature_dim) state_dict = torch.load(model_path) - # state_dict のキーを修正 - new_state_dict = {} - for key in state_dict: - new_key = "model." + key - new_state_dict[new_key] = state_dict[key] - model.load_state_dict(new_state_dict) + # # state_dict のキーを修正 + # new_state_dict = {} + # for key in state_dict: + # new_key = "model." + key + # new_state_dict[new_key] = state_dict[key] + model.load_state_dict(state_dict['state_dict']) model.eval() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -58,7 +58,9 @@ x = load_data(device) # アップロード画像の前処理 def preprocess_uploaded_image(uploaded_image, image_size): - uploaded_image = Image.fromarray(uploaded_image) + # ndarrayの場合はPILイメージに変換 + if type(uploaded_image) == np.ndarray: + uploaded_image = Image.fromarray(uploaded_image) uploaded_image = uploaded_image.convert("RGB") uploaded_image = uploaded_image.resize((image_size, image_size)) uploaded_image = np.array(uploaded_image).transpose(2, 0, 1) / 255.0 @@ -66,8 +68,17 @@ def preprocess_uploaded_image(uploaded_image, image_size): return uploaded_image # ヒートマップの生成関数 -@spaces.GPU +# @spaces.GPU def get_heatmaps(source_num, x_coords, y_coords, uploaded_image): + if type(uploaded_image) == str: + uploaded_image = Image.open(uploaded_image) + if type(source_num) == str: + source_num = int(source_num) + if type(x_coords) == str: + x_coords = int(x_coords) + if type(y_coords) == str: + y_coords = int(y_coords) + with torch.no_grad(): dec5, _ = model(x) img = x @@ -101,7 +112,7 @@ def get_heatmaps(source_num, x_coords, y_coords, uploaded_image): source_map = norm_batch_distance_map[source_num] target_map = norm_batch_distance_map[target_num] - alpha = 0.8 + alpha = 0.7 blended_source = (1 - alpha) * img[source_num] + alpha * torch.cat(((norm_batch_distance_map[source_num] / norm_batch_distance_map[source_num].max()).unsqueeze(0), torch.zeros(2, image_size, image_size, device=device))) blended_target = (1 - alpha) * img[target_num] + alpha * torch.cat(((norm_batch_distance_map[target_num] / norm_batch_distance_map[target_num].max()).unsqueeze(0), torch.zeros(2, image_size, image_size, device=device))) @@ -154,6 +165,7 @@ async () => { console.log(files); if (files && files.length > 0) { console.log("File selected"); + document.querySelector("#input_file_button").style.display = "none"; document.querySelector("#crop_view").style.display = "block"; document.querySelector("#crop_button").style.display = "block"; const url = URL.createObjectURL(files[0]); @@ -183,6 +195,7 @@ async () => { document.getElementById("crop_view").style.display = "none"; document.getElementById("crop_button").style.display = "none"; + document.querySelector("#input_file_button").style.display = "block"; cropper.destroy(); } @@ -194,6 +207,16 @@ async () => { """ with gr.Blocks() as demo: + # title + gr.Markdown("# TripletGeoEncoder Feature Map Visualization") + # description + gr.Markdown("This demo visualizes the feature maps of a TripletGeoEncoder trained on the CelebA dataset using self-supervised learning without annotations from only 1000 images. " + "The feature maps are visualized as heatmaps, where the source map shows the distance of each pixel in the source image to the selected pixel, and the target map shows the distance of each pixel in the target image to the selected pixel. " + + "The blended source and target images show the source and target images with the source and target maps overlaid, respectively. " + + "For further information, please contact me on X (formerly Twitter): @Yeq6X.") + with gr.Row(): with gr.Column(): source_num = gr.Slider(0, batch_size - 1, step=1, label="Source Image Index") @@ -202,20 +225,30 @@ with gr.Blocks() as demo: # GradioのFileコンポーネントでファイル選択ボタンを追加 gr.HTML('') - input_file_button = gr.Button("画像を選択", elem_id="input_file_button") + input_file_button = gr.Button("Upload Image and Crop", elem_id="input_file_button", variant="primary") + crop_button = gr.Button("Crop", elem_id="crop_button", variant="primary") # 画像を表示するためのHTML画像タグをGradioで表示 gr.HTML('') # Gradioのボタンコンポーネントを追加し、IDを付与 - crop_button = gr.Button("クロップ", elem_id="crop_button", variant="primary") # クロップされた画像データのテキストボックス(Base64データ) cropped_image_data = gr.Textbox(visible=False, elem_id="cropped_image_data") - input_image = gr.Image(label="Cropped Image", interactive=False) + input_image = gr.Image(label="Cropped Image", elem_id="input_image") # cropped_image_dataが更新されたらprocess_imageを呼び出す cropped_image_data.change(process_image, inputs=cropped_image_data, outputs=input_image) + # examples + gr.Markdown("# Examples") + gr.Examples( + examples=[ + ["0", "50", "50", "resources/examples/2488.jpg"], + ["0", "50", "50", "resources/examples/2899.jpg"] + ], + inputs=[source_num, x_coords, y_coords, input_image], + ) with gr.Column(): output_plot = gr.Plot() + # Gradioインターフェースの代わり source_num.change(get_heatmaps, inputs=[source_num, x_coords, y_coords, input_image], outputs=output_plot) x_coords.change(get_heatmaps, inputs=[source_num, x_coords, y_coords, input_image], outputs=output_plot) @@ -223,7 +256,7 @@ with gr.Blocks() as demo: input_image.change(get_heatmaps, inputs=[source_num, x_coords, y_coords, input_image], outputs=output_plot) # JavaScriptコードをロード - demo.load(None, None, None, js=scripts) - + demo.load(None, None, None, js=scripts) + demo.launch() diff --git a/resources/examples/2488.jpg b/resources/examples/2488.jpg new file mode 100644 index 0000000000000000000000000000000000000000..754e7bf0a370f0fe3974858ca385e981bbf08d37 Binary files /dev/null and b/resources/examples/2488.jpg differ diff --git a/resources/examples/2899.jpg b/resources/examples/2899.jpg new file mode 100644 index 0000000000000000000000000000000000000000..be036cce96c26c03d8611d7ad43cc92e14ddfc3f Binary files /dev/null and b/resources/examples/2899.jpg differ diff --git a/resources/trainB/0000.jpg b/resources/trainB/0000.jpg new file mode 100644 index 0000000000000000000000000000000000000000..d1c4bd0506078d7cb49c488c4e55238079e2f2d1 Binary files /dev/null and b/resources/trainB/0000.jpg differ diff --git a/resources/trainB/0001.jpg b/resources/trainB/0001.jpg new file mode 100644 index 0000000000000000000000000000000000000000..e54fc3ef2b6d52a986544fbbeeec217e723f0f30 Binary files /dev/null and b/resources/trainB/0001.jpg differ diff --git a/resources/trainB/0002.jpg b/resources/trainB/0002.jpg new file mode 100644 index 0000000000000000000000000000000000000000..95716ec26e65ea2324b5ebb323bb118203cf260e Binary files /dev/null and b/resources/trainB/0002.jpg differ diff --git a/resources/trainB/0003.jpg b/resources/trainB/0003.jpg new file mode 100644 index 0000000000000000000000000000000000000000..caff69928a18a74725d56918d8de9849eacb09c4 Binary files /dev/null and b/resources/trainB/0003.jpg differ diff --git a/resources/trainB/0004.jpg b/resources/trainB/0004.jpg new file mode 100644 index 0000000000000000000000000000000000000000..7f42419469490865a0d168eadbd8750e2ec72cd7 Binary files /dev/null and b/resources/trainB/0004.jpg differ diff --git a/resources/trainB/0005.jpg b/resources/trainB/0005.jpg new file mode 100644 index 0000000000000000000000000000000000000000..d643652082fc9ecf9af4c79618f27e78cd871697 Binary files /dev/null and b/resources/trainB/0005.jpg differ diff --git a/resources/trainB/0006.jpg b/resources/trainB/0006.jpg new file mode 100644 index 0000000000000000000000000000000000000000..3697f3410c390429bd0e959b1c628a0208f3578a Binary files /dev/null and b/resources/trainB/0006.jpg differ diff --git a/resources/trainB/0007.jpg b/resources/trainB/0007.jpg new file mode 100644 index 0000000000000000000000000000000000000000..21e2126319fdd10d4ae5afe4785a7f10514d79ad Binary files /dev/null and b/resources/trainB/0007.jpg differ diff --git a/resources/trainB/0008.jpg b/resources/trainB/0008.jpg new file mode 100644 index 0000000000000000000000000000000000000000..220bd17b9d858121ef414afff7190e68066ca852 Binary files /dev/null and b/resources/trainB/0008.jpg differ diff --git a/resources/trainB/0009.jpg b/resources/trainB/0009.jpg new file mode 100644 index 0000000000000000000000000000000000000000..611c14f8a0c7eb0206a530aed288051ddd15d2b2 Binary files /dev/null and b/resources/trainB/0009.jpg differ diff --git a/resources/trainB/0010.jpg b/resources/trainB/0010.jpg new file mode 100644 index 0000000000000000000000000000000000000000..17e2894ba8662edca8b74422fa2c7391b4800091 Binary files /dev/null and b/resources/trainB/0010.jpg differ diff --git a/resources/trainB/0011.jpg b/resources/trainB/0011.jpg new file mode 100644 index 0000000000000000000000000000000000000000..b826433e78e058f87f4a99b1a6a82f08152d10ca Binary files /dev/null and b/resources/trainB/0011.jpg differ diff --git a/resources/trainB/0012.jpg b/resources/trainB/0012.jpg new file mode 100644 index 0000000000000000000000000000000000000000..374c92c88167f67e94e7907561f82b4658717467 Binary files /dev/null and b/resources/trainB/0012.jpg differ diff --git a/resources/trainB/0013.jpg b/resources/trainB/0013.jpg new file mode 100644 index 0000000000000000000000000000000000000000..aaa8c05ddf8df180c9742bb88e09808cfef03f1a Binary files /dev/null and b/resources/trainB/0013.jpg differ diff --git a/resources/trainB/0014.jpg b/resources/trainB/0014.jpg new file mode 100644 index 0000000000000000000000000000000000000000..3145de9318bfb6b99862236812a4a87769be82c9 Binary files /dev/null and b/resources/trainB/0014.jpg differ diff --git a/resources/trainB/0015.jpg b/resources/trainB/0015.jpg new file mode 100644 index 0000000000000000000000000000000000000000..d44765d66a3dc3d6e277c9c299b2166bdf2f15af Binary files /dev/null and b/resources/trainB/0015.jpg differ diff --git a/resources/trainB/0016.jpg b/resources/trainB/0016.jpg new file mode 100644 index 0000000000000000000000000000000000000000..a6986cc8b22f36058a16b6006e0629ee6704759c Binary files /dev/null and b/resources/trainB/0016.jpg differ diff --git a/resources/trainB/0017.jpg b/resources/trainB/0017.jpg new file mode 100644 index 0000000000000000000000000000000000000000..112430d387b2cf158ea33d6757ec97919ec79084 Binary files /dev/null and b/resources/trainB/0017.jpg differ diff --git a/resources/trainB/0018.jpg b/resources/trainB/0018.jpg new file mode 100644 index 0000000000000000000000000000000000000000..2e676a22ad2ad7c6ae8cb9f78ffc997ae833606d Binary files /dev/null and b/resources/trainB/0018.jpg differ diff --git a/resources/trainB/0019.jpg b/resources/trainB/0019.jpg new file mode 100644 index 0000000000000000000000000000000000000000..66d73897db8e73e47b2f93125c7e699d47ae6efa Binary files /dev/null and b/resources/trainB/0019.jpg differ diff --git a/resources/trainB/0020.jpg b/resources/trainB/0020.jpg new file mode 100644 index 0000000000000000000000000000000000000000..7f8ea86893e1b1cd95a72a2d1924f91a3a5435f8 Binary files /dev/null and b/resources/trainB/0020.jpg differ diff --git a/resources/trainB/0021.jpg b/resources/trainB/0021.jpg new file mode 100644 index 0000000000000000000000000000000000000000..4e79015be10b45336ad6ca3afab3e6032a176216 Binary files /dev/null and b/resources/trainB/0021.jpg differ diff --git a/resources/trainB/0022.jpg b/resources/trainB/0022.jpg new file mode 100644 index 0000000000000000000000000000000000000000..6233d663aa29e9bef2b65b44dc6f7491781fd33a Binary files /dev/null and b/resources/trainB/0022.jpg differ diff --git a/resources/trainB/0023.jpg b/resources/trainB/0023.jpg new file mode 100644 index 0000000000000000000000000000000000000000..36eee9e9614ea0650c42e3ea8958e959f14a1062 Binary files /dev/null and b/resources/trainB/0023.jpg differ diff --git a/resources/trainB/0024.jpg b/resources/trainB/0024.jpg new file mode 100644 index 0000000000000000000000000000000000000000..328ccc567a5306f7868e7857ce5f963eef192889 Binary files /dev/null and b/resources/trainB/0024.jpg differ diff --git a/resources/trainB/0025.jpg b/resources/trainB/0025.jpg new file mode 100644 index 0000000000000000000000000000000000000000..57474a2e9e5553195f23b78a8cbb8aa66f28b96c Binary files /dev/null and b/resources/trainB/0025.jpg differ diff --git a/resources/trainB/0026.jpg b/resources/trainB/0026.jpg new file mode 100644 index 0000000000000000000000000000000000000000..2e0025cee4156aef800bb804f1b22a1cea573b51 Binary files /dev/null and b/resources/trainB/0026.jpg differ diff --git a/resources/trainB/0027.jpg b/resources/trainB/0027.jpg new file mode 100644 index 0000000000000000000000000000000000000000..69d7fd242ab260fec412f4d4202f4c2263086ed5 Binary files /dev/null and b/resources/trainB/0027.jpg differ diff --git a/resources/trainB/0028.jpg b/resources/trainB/0028.jpg new file mode 100644 index 0000000000000000000000000000000000000000..a95d9a117d3778af067006e78995445401fec764 Binary files /dev/null and b/resources/trainB/0028.jpg differ diff --git a/resources/trainB/0029.jpg b/resources/trainB/0029.jpg new file mode 100644 index 0000000000000000000000000000000000000000..81d3ae24dbceca43dbc90ce112e1cda66d112be3 Binary files /dev/null and b/resources/trainB/0029.jpg differ diff --git a/resources/trainB/0030.jpg b/resources/trainB/0030.jpg new file mode 100644 index 0000000000000000000000000000000000000000..a52c579d28063d28d1bbfe807b0a81fb1c538a78 Binary files /dev/null and b/resources/trainB/0030.jpg differ diff --git a/resources/trainB/0031.jpg b/resources/trainB/0031.jpg new file mode 100644 index 0000000000000000000000000000000000000000..5d3c2023ae9b3244a3e49db6c32e1ca84f03b035 Binary files /dev/null and b/resources/trainB/0031.jpg differ diff --git a/resources/trainB/0032.jpg b/resources/trainB/0032.jpg new file mode 100644 index 0000000000000000000000000000000000000000..10a610beb2ed6bcc1e9e27fa9a7ffca090d92eee Binary files /dev/null and b/resources/trainB/0032.jpg differ diff --git a/resources/trainB/0033.jpg b/resources/trainB/0033.jpg new file mode 100644 index 0000000000000000000000000000000000000000..2b53ec3b6e29188ee52aa466506571f295781318 Binary files /dev/null and b/resources/trainB/0033.jpg differ diff --git a/resources/trainB/0034.jpg b/resources/trainB/0034.jpg new file mode 100644 index 0000000000000000000000000000000000000000..d4bd665a61e4d3da65dc520a7de8e7d95df59653 Binary files /dev/null and b/resources/trainB/0034.jpg differ diff --git a/resources/trainB/0035.jpg b/resources/trainB/0035.jpg new file mode 100644 index 0000000000000000000000000000000000000000..1e48562019f931ff93954a2876ba2ed3c9f94d56 Binary files /dev/null and b/resources/trainB/0035.jpg differ diff --git a/resources/trainB/0036.jpg b/resources/trainB/0036.jpg new file mode 100644 index 0000000000000000000000000000000000000000..48cbbba2e3e879289549117048dbc2637939741d Binary files /dev/null and b/resources/trainB/0036.jpg differ diff --git a/resources/trainB/0037.jpg b/resources/trainB/0037.jpg new file mode 100644 index 0000000000000000000000000000000000000000..e93a9fab47803532313c3c5ac576a0c62b6e8f2e Binary files /dev/null and b/resources/trainB/0037.jpg differ diff --git a/resources/trainB/0038.jpg b/resources/trainB/0038.jpg new file mode 100644 index 0000000000000000000000000000000000000000..c533679fa7ea2b8a8cf5f8da38055c1fb505dc9a Binary files /dev/null and b/resources/trainB/0038.jpg differ diff --git a/resources/trainB/0039.jpg b/resources/trainB/0039.jpg new file mode 100644 index 0000000000000000000000000000000000000000..534fa849da42f87f941fe0fa0094db8c5d991198 Binary files /dev/null and b/resources/trainB/0039.jpg differ diff --git a/resources/trainB/0040.jpg b/resources/trainB/0040.jpg new file mode 100644 index 0000000000000000000000000000000000000000..ebde959a9743d0fd4b9d01397e0bd9927b6712ef Binary files /dev/null and b/resources/trainB/0040.jpg differ diff --git a/resources/trainB/0041.jpg b/resources/trainB/0041.jpg new file mode 100644 index 0000000000000000000000000000000000000000..e891a3a472517d2d63318f6ccf48fdceff6ccb4c Binary files /dev/null and b/resources/trainB/0041.jpg differ diff --git a/resources/trainB/0042.jpg b/resources/trainB/0042.jpg new file mode 100644 index 0000000000000000000000000000000000000000..cbc5eea650e28acaee079344d1fbf2a72b6b3911 Binary files /dev/null and b/resources/trainB/0042.jpg differ diff --git a/resources/trainB/0043.jpg b/resources/trainB/0043.jpg new file mode 100644 index 0000000000000000000000000000000000000000..3a6cb99f59b694dad8ad2ddd828ae561200c4f02 Binary files /dev/null and b/resources/trainB/0043.jpg differ diff --git a/resources/trainB/0044.jpg b/resources/trainB/0044.jpg new file mode 100644 index 0000000000000000000000000000000000000000..70c9f10912c4f8745175f700fc4b5e8dd9ecbdc9 Binary files /dev/null and b/resources/trainB/0044.jpg differ diff --git a/resources/trainB/0045.jpg b/resources/trainB/0045.jpg new file mode 100644 index 0000000000000000000000000000000000000000..22d39ab7a73c14fde203bbe8f77b736cf4ddc2f2 Binary files /dev/null and b/resources/trainB/0045.jpg differ diff --git a/resources/trainB/0046.jpg b/resources/trainB/0046.jpg new file mode 100644 index 0000000000000000000000000000000000000000..ef5dbba6a1bc7ade18837c344ff069b73ee25547 Binary files /dev/null and b/resources/trainB/0046.jpg differ diff --git a/resources/trainB/0047.jpg b/resources/trainB/0047.jpg new file mode 100644 index 0000000000000000000000000000000000000000..94257ea02231fc3cdb19423ff8d09b5476b110c8 Binary files /dev/null and b/resources/trainB/0047.jpg differ diff --git a/resources/trainB/0048.jpg b/resources/trainB/0048.jpg new file mode 100644 index 0000000000000000000000000000000000000000..2d412e6cf8a4ab655925a22bfc938ad8881a4414 Binary files /dev/null and b/resources/trainB/0048.jpg differ diff --git a/resources/trainB/0049.jpg b/resources/trainB/0049.jpg new file mode 100644 index 0000000000000000000000000000000000000000..f3828f979e802c4c92fb1f084205b391739872fe Binary files /dev/null and b/resources/trainB/0049.jpg differ diff --git a/resources/trainB/0050.jpg b/resources/trainB/0050.jpg new file mode 100644 index 0000000000000000000000000000000000000000..d59bd7068eb6a727099cf8c751511cd6d542b4a6 Binary files /dev/null and b/resources/trainB/0050.jpg differ diff --git a/resources/trainB/0051.jpg b/resources/trainB/0051.jpg new file mode 100644 index 0000000000000000000000000000000000000000..f9a9252d03acefa9541296919e450f250c704579 Binary files /dev/null and b/resources/trainB/0051.jpg differ diff --git a/resources/trainB/0052.jpg b/resources/trainB/0052.jpg new file mode 100644 index 0000000000000000000000000000000000000000..6e674cf6b5eefae0e0160c125b6057e671e0c5df Binary files /dev/null and b/resources/trainB/0052.jpg differ diff --git a/resources/trainB/0053.jpg b/resources/trainB/0053.jpg new file mode 100644 index 0000000000000000000000000000000000000000..9b699b3f0754ab9cd4679e84cd9b6b26a04b34d3 Binary files /dev/null and b/resources/trainB/0053.jpg differ diff --git a/resources/trainB/0054.jpg b/resources/trainB/0054.jpg new file mode 100644 index 0000000000000000000000000000000000000000..c33ff626d14e8814422eddbbbf8fc1c9d3c8e738 Binary files /dev/null and b/resources/trainB/0054.jpg differ diff --git a/resources/trainB/0055.jpg b/resources/trainB/0055.jpg new file mode 100644 index 0000000000000000000000000000000000000000..24fc94cd2f60c84185cc31ae394a09bf00b9a727 Binary files /dev/null and b/resources/trainB/0055.jpg differ diff --git a/resources/trainB/0056.jpg b/resources/trainB/0056.jpg new file mode 100644 index 0000000000000000000000000000000000000000..53c951674bcdb9de94ec98e883a7bc21b6e82c3b Binary files /dev/null and b/resources/trainB/0056.jpg differ diff --git a/resources/trainB/0057.jpg b/resources/trainB/0057.jpg new file mode 100644 index 0000000000000000000000000000000000000000..0dd7fe2b56c0024e61a26a26333ac2ba13ff6102 Binary files /dev/null and b/resources/trainB/0057.jpg differ diff --git a/resources/trainB/0058.jpg b/resources/trainB/0058.jpg new file mode 100644 index 0000000000000000000000000000000000000000..8c86ea7c2605a9ddecad7cfda2ae22f779b8fd73 Binary files /dev/null and b/resources/trainB/0058.jpg differ diff --git a/resources/trainB/0059.jpg b/resources/trainB/0059.jpg new file mode 100644 index 0000000000000000000000000000000000000000..17b5ff64d45d97e3cbfde6cce320b6372ec7fc7f Binary files /dev/null and b/resources/trainB/0059.jpg differ diff --git a/resources/trainB/0060.jpg b/resources/trainB/0060.jpg new file mode 100644 index 0000000000000000000000000000000000000000..177d986045d0e47c4844424debe0c93288aa61a7 Binary files /dev/null and b/resources/trainB/0060.jpg differ diff --git a/resources/trainB/0061.jpg b/resources/trainB/0061.jpg new file mode 100644 index 0000000000000000000000000000000000000000..d4f6990c1ebe5d843e901360ac341baf7d087a05 Binary files /dev/null and b/resources/trainB/0061.jpg differ diff --git a/resources/trainB/0062.jpg b/resources/trainB/0062.jpg new file mode 100644 index 0000000000000000000000000000000000000000..d8eb612a877924f1f56199b9d61aee50be21f26c Binary files /dev/null and b/resources/trainB/0062.jpg differ diff --git a/resources/trainB/0063.jpg b/resources/trainB/0063.jpg new file mode 100644 index 0000000000000000000000000000000000000000..6bacbf61260a3c77685fa9f6ca95e244b58ce97e Binary files /dev/null and b/resources/trainB/0063.jpg differ