Spaces:
Sleeping
Sleeping
File size: 3,711 Bytes
c0fe349 738801e c0fe349 738801e c0fe349 738801e c0fe349 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
{
"cells": [
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"c:\\Users\\hiroga\\miniconda3\\envs\\pokemon-pal\\Lib\\site-packages\\torchvision\\transforms\\functional.py:1603: UserWarning: The default value of the antialias parameter of all the resizing transforms (Resize(), RandomResizedCrop(), etc.) will change from None to True in v0.17, in order to be consistent across the PIL and Tensor backends. To suppress this warning, directly pass antialias=True (recommended, future default), antialias=None (current default, which means False for Tensors and True for PIL), or antialias=False (only works on Tensors - PIL will still use antialiasing). This also applies if you are using the inference transforms from the models weights: update the call to weights.transforms(antialias=True).\n",
" warnings.warn(\n"
]
}
],
"source": [
"import torch\n",
"from carvekit.api.high import HiInterface\n",
"\n",
"# Check doc strings for more information\n",
"interface = HiInterface(object_type=\"object\", # Can be \"object\" or \"hairs-like\".\n",
" batch_size_seg=5,\n",
" batch_size_matting=1,\n",
" device='cuda' if torch.cuda.is_available() else 'cpu',\n",
" seg_mask_size=640, # Use 640 for Tracer B7 and 320 for U2Net\n",
" matting_mask_size=2048,\n",
" trimap_prob_threshold=231,\n",
" trimap_dilation=30,\n",
" trimap_erosion_iters=5,\n",
" fp16=False)\n",
"import os\n",
"\n",
"# input_dir = \"../data/raw\"\n",
"# output_dir = \"../data/nobg\"\n",
"input_dir = \"../data/raw/paldex.io\"\n",
"output_dir = \"../data/nobg/paldex.io\"\n",
"\n",
"# Create output directory if it doesn't exist\n",
"os.makedirs(output_dir, exist_ok=True)\n",
"\n",
"# Loop over all files and subdirectories in the input directory\n",
"for root, dirs, files in os.walk(input_dir):\n",
" for filename in files:\n",
" # Construct full file path\n",
" file_path = os.path.join(root, filename)\n",
" \n",
" # Process the image and remove the background\n",
" images_without_background = interface([file_path])\n",
" image_wo_bg = images_without_background[0]\n",
" \n",
" # Create output subdirectory if it doesn't exist\n",
" output_subdir = os.path.join(output_dir, os.path.relpath(root, input_dir))\n",
" os.makedirs(output_subdir, exist_ok=True)\n",
" \n",
" # Save the processed image to the output directory\n",
" # Since the image format is RGBA, we save it as PNG\n",
" filename = os.path.splitext(filename)[0] + \".png\"\n",
" output_file_path = os.path.join(output_subdir, filename)\n",
" image_wo_bg.save(output_file_path)\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "pokemon-pal",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.7"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
|