Sample usage.
This commit is contained in:
parent
d951e2369d
commit
b485d4e838
|
|
@ -0,0 +1,175 @@
|
||||||
|
#!/bin/bash
|
||||||
|
set -e
|
||||||
|
|
||||||
|
install_opencv () {
|
||||||
|
# Check if the file /proc/device-tree/model exists
|
||||||
|
if [ -e "/proc/device-tree/model" ]; then
|
||||||
|
# Read the model information from /proc/device-tree/model and remove null bytes
|
||||||
|
model=$(tr -d '\0' < /proc/device-tree/model)
|
||||||
|
# Check if the model information contains "Jetson Nano Orion"
|
||||||
|
echo ""
|
||||||
|
if [[ $model == *"Orin"* ]]; then
|
||||||
|
echo "Detecting a Jetson Nano Orin."
|
||||||
|
NO_JOB=4
|
||||||
|
ARCH=8.7
|
||||||
|
PTX="sm_87"
|
||||||
|
elif [[ $model == *"Jetson Nano"* ]]; then
|
||||||
|
echo "Detecting a regular Jetson Nano."
|
||||||
|
ARCH=5.3
|
||||||
|
PTX="sm_53"
|
||||||
|
# Use "-j 4" only swap space is larger than 5.5GB
|
||||||
|
FREE_MEM="$(free -m | awk '/^Swap/ {print $2}')"
|
||||||
|
if [[ "FREE_MEM" -gt "5500" ]]; then
|
||||||
|
NO_JOB=4
|
||||||
|
else
|
||||||
|
echo "Due to limited swap, make only uses 1 core"
|
||||||
|
NO_JOB=1
|
||||||
|
fi
|
||||||
|
else
|
||||||
|
echo "Unable to determine the Jetson Nano model."
|
||||||
|
exit 1
|
||||||
|
fi
|
||||||
|
echo ""
|
||||||
|
else
|
||||||
|
echo "Error: /proc/device-tree/model not found. Are you sure this is a Jetson Nano?"
|
||||||
|
exit 1
|
||||||
|
fi
|
||||||
|
|
||||||
|
echo "Installing OpenCV 4.9.0 on your Nano"
|
||||||
|
echo "It will take 3.5 hours !"
|
||||||
|
|
||||||
|
# reveal the CUDA location
|
||||||
|
cd ~
|
||||||
|
sudo sh -c "echo '/usr/local/cuda/lib64' >> /etc/ld.so.conf.d/nvidia-tegra.conf"
|
||||||
|
sudo ldconfig
|
||||||
|
|
||||||
|
# install the Jetson Nano dependencies first
|
||||||
|
if [[ $model == *"Jetson Nano"* ]]; then
|
||||||
|
sudo apt-get install -y build-essential git unzip pkg-config zlib1g-dev
|
||||||
|
sudo apt-get install -y python3-dev python3-numpy python3-pip
|
||||||
|
sudo apt-get install -y gstreamer1.0-tools libgstreamer-plugins-base1.0-dev
|
||||||
|
sudo apt-get install -y libgstreamer-plugins-good1.0-dev
|
||||||
|
sudo apt-get install -y libtbb2 libgtk-3-dev v4l2ucp libxine2-dev
|
||||||
|
fi
|
||||||
|
|
||||||
|
if [ -f /etc/os-release ]; then
|
||||||
|
# Source the /etc/os-release file to get variables
|
||||||
|
. /etc/os-release
|
||||||
|
# Extract the major version number from VERSION_ID
|
||||||
|
VERSION_MAJOR=$(echo "$VERSION_ID" | cut -d'.' -f1)
|
||||||
|
# Check if the extracted major version is 22 or earlier
|
||||||
|
if [ "$VERSION_MAJOR" = "22" ]; then
|
||||||
|
sudo apt-get install -y libswresample-dev libdc1394-dev
|
||||||
|
else
|
||||||
|
sudo apt-get install -y libavresample-dev libdc1394-22-dev
|
||||||
|
fi
|
||||||
|
else
|
||||||
|
sudo apt-get install -y libavresample-dev libdc1394-22-dev
|
||||||
|
fi
|
||||||
|
|
||||||
|
# install the common dependencies
|
||||||
|
sudo apt-get install -y cmake
|
||||||
|
sudo apt-get install -y libjpeg-dev libjpeg8-dev libjpeg-turbo8-dev
|
||||||
|
sudo apt-get install -y libpng-dev libtiff-dev libglew-dev
|
||||||
|
sudo apt-get install -y libavcodec-dev libavformat-dev libswscale-dev
|
||||||
|
sudo apt-get install -y libgtk2.0-dev libgtk-3-dev libcanberra-gtk*
|
||||||
|
sudo apt-get install -y libxvidcore-dev libx264-dev
|
||||||
|
sudo apt-get install -y libtbb-dev libxine2-dev
|
||||||
|
sudo apt-get install -y libv4l-dev v4l-utils qv4l2
|
||||||
|
sudo apt-get install -y libtesseract-dev libpostproc-dev
|
||||||
|
sudo apt-get install -y libvorbis-dev
|
||||||
|
sudo apt-get install -y libfaac-dev libmp3lame-dev libtheora-dev
|
||||||
|
sudo apt-get install -y libopencore-amrnb-dev libopencore-amrwb-dev
|
||||||
|
sudo apt-get install -y libopenblas-dev libatlas-base-dev libblas-dev
|
||||||
|
sudo apt-get install -y liblapack-dev liblapacke-dev libeigen3-dev gfortran
|
||||||
|
sudo apt-get install -y libhdf5-dev libprotobuf-dev protobuf-compiler
|
||||||
|
sudo apt-get install -y libgoogle-glog-dev libgflags-dev
|
||||||
|
|
||||||
|
# remove old versions or previous builds
|
||||||
|
cd ~
|
||||||
|
sudo rm -rf opencv*
|
||||||
|
# download the latest version
|
||||||
|
git clone --depth=1 https://github.com/opencv/opencv.git
|
||||||
|
git clone --depth=1 https://github.com/opencv/opencv_contrib.git
|
||||||
|
|
||||||
|
# set install dir
|
||||||
|
cd ~/opencv
|
||||||
|
mkdir build
|
||||||
|
cd build
|
||||||
|
|
||||||
|
# run cmake
|
||||||
|
cmake -D CMAKE_BUILD_TYPE=RELEASE \
|
||||||
|
-D CMAKE_INSTALL_PREFIX=/usr \
|
||||||
|
-D OPENCV_EXTRA_MODULES_PATH=~/opencv_contrib/modules \
|
||||||
|
-D EIGEN_INCLUDE_PATH=/usr/include/eigen3 \
|
||||||
|
-D WITH_OPENCL=OFF \
|
||||||
|
-D CUDA_ARCH_BIN=${ARCH} \
|
||||||
|
-D CUDA_ARCH_PTX=${PTX} \
|
||||||
|
-D WITH_CUDA=ON \
|
||||||
|
-D WITH_CUDNN=ON \
|
||||||
|
-D WITH_CUBLAS=ON \
|
||||||
|
-D ENABLE_FAST_MATH=ON \
|
||||||
|
-D CUDA_FAST_MATH=ON \
|
||||||
|
-D OPENCV_DNN_CUDA=ON \
|
||||||
|
-D ENABLE_NEON=ON \
|
||||||
|
-D WITH_QT=OFF \
|
||||||
|
-D WITH_OPENMP=ON \
|
||||||
|
-D BUILD_TIFF=ON \
|
||||||
|
-D WITH_FFMPEG=ON \
|
||||||
|
-D WITH_GSTREAMER=ON \
|
||||||
|
-D WITH_TBB=ON \
|
||||||
|
-D BUILD_TBB=ON \
|
||||||
|
-D BUILD_TESTS=OFF \
|
||||||
|
-D WITH_EIGEN=ON \
|
||||||
|
-D WITH_V4L=ON \
|
||||||
|
-D WITH_LIBV4L=ON \
|
||||||
|
-D WITH_PROTOBUF=ON \
|
||||||
|
-D INSTALL_C_EXAMPLES=OFF \
|
||||||
|
-D INSTALL_PYTHON_EXAMPLES=OFF \
|
||||||
|
-D PYTHON3_EXECUTABLE=$(which python3) \
|
||||||
|
-D PYTHON3_INCLUDE_DIR=$(python3 -c "from distutils.sysconfig import get_python_inc; print(get_python_inc())") \
|
||||||
|
-D PYTHON3_LIBRARY=$(python3 -c "import distutils.sysconfig as sysconfig; print(sysconfig.get_config_var('LIBDIR'))") \
|
||||||
|
-D PYTHON3_PACKAGES_PATH=$(python3 -c "import site; print(site.getsitepackages()[0])") \
|
||||||
|
-D OPENCV_GENERATE_PKGCONFIG=ON \
|
||||||
|
-D BUILD_EXAMPLES=OFF \
|
||||||
|
-D CMAKE_CXX_FLAGS="-march=native -mtune=native" \
|
||||||
|
-D CMAKE_C_FLAGS="-march=native -mtune=native" ..
|
||||||
|
|
||||||
|
make -j ${NO_JOB}
|
||||||
|
|
||||||
|
directory="/usr/include/opencv4/opencv2"
|
||||||
|
if [ -d "$directory" ]; then
|
||||||
|
# Directory exists, so delete it
|
||||||
|
sudo rm -rf "$directory"
|
||||||
|
fi
|
||||||
|
|
||||||
|
sudo make install
|
||||||
|
sudo ldconfig
|
||||||
|
|
||||||
|
# cleaning (frees 320 MB)
|
||||||
|
make clean
|
||||||
|
sudo apt-get update
|
||||||
|
|
||||||
|
echo "Congratulations!"
|
||||||
|
echo "You've successfully installed OpenCV 4.9.0 on your Nano with Python 3 bindings."
|
||||||
|
}
|
||||||
|
|
||||||
|
cd ~
|
||||||
|
|
||||||
|
if [ -d ~/opencv/build ]; then
|
||||||
|
echo " "
|
||||||
|
echo "You have a directory ~/opencv/build on your disk."
|
||||||
|
echo "Continuing the installation will replace this folder."
|
||||||
|
echo " "
|
||||||
|
|
||||||
|
printf "Do you wish to continue (Y/n)?"
|
||||||
|
read answer
|
||||||
|
|
||||||
|
if [ "$answer" != "${answer#[Nn]}" ] ;then
|
||||||
|
echo "Leaving without installing OpenCV"
|
||||||
|
else
|
||||||
|
install_opencv
|
||||||
|
fi
|
||||||
|
else
|
||||||
|
install_opencv
|
||||||
|
fi
|
||||||
|
|
@ -0,0 +1,26 @@
|
||||||
|
https://forums.developer.nvidia.com/t/which-tracker-to-use-to-do-ai-detection-1st-frame-tracking-other-frames/178305/6
|
||||||
|
|
||||||
|
1. sudo apt-get install python3-dev
|
||||||
|
|
||||||
|
2. install dlib for cuda
|
||||||
|
wget http://dlib.net/files/dlib-19.22.tar.bz2
|
||||||
|
tar jxvf dlib-19.22.tar.bz2
|
||||||
|
cd dlib-19.22/
|
||||||
|
mkdir build
|
||||||
|
cd build/
|
||||||
|
cmake ..
|
||||||
|
cmake --build .
|
||||||
|
cd ../
|
||||||
|
sudo python3 setup.py install
|
||||||
|
|
||||||
|
3.
|
||||||
|
python3
|
||||||
|
|
||||||
|
import dlib
|
||||||
|
dlib.cuda.get_num_devices()
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
For opencv with virtualenv:
|
||||||
|
cp -r /usr/local/lib/python3.6/dist-packages/cv2 .
|
||||||
|
cp -r /usr/lib/python3/dist-packages/numpy .
|
||||||
|
|
@ -0,0 +1,100 @@
|
||||||
|
import cv2
|
||||||
|
import dlib
|
||||||
|
import time # To measure time for FPS calculation
|
||||||
|
|
||||||
|
# GStreamer pipeline to acquire video from the CSI camera (Jetson Nano)
|
||||||
|
def gstreamer_pipeline(
|
||||||
|
sensor_id=0,
|
||||||
|
capture_width=1280,
|
||||||
|
capture_height=720,
|
||||||
|
display_width=640,
|
||||||
|
display_height=360,
|
||||||
|
framerate=15,
|
||||||
|
flip_method=0,
|
||||||
|
):
|
||||||
|
return (
|
||||||
|
f"nvarguscamerasrc sensor-id={sensor_id} ! "
|
||||||
|
f"video/x-raw(memory:NVMM), width=(int){capture_width}, height=(int){capture_height}, "
|
||||||
|
f"format=(string)NV12, framerate=(fraction){framerate}/1 ! "
|
||||||
|
f"nvvidconv flip-method={flip_method} ! "
|
||||||
|
f"video/x-raw, width=(int){display_width}, height=(int){display_height}, format=(string)BGRx ! "
|
||||||
|
f"videoconvert ! video/x-raw, format=(string)BGR ! appsink"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Capture from the camera using OpenCV and start tracking
|
||||||
|
cap = cv2.VideoCapture(gstreamer_pipeline(), cv2.CAP_GSTREAMER)
|
||||||
|
|
||||||
|
if not cap.isOpened():
|
||||||
|
print("Unable to open camera")
|
||||||
|
else:
|
||||||
|
# Read the first frame to allow user to draw a bounding box
|
||||||
|
ret, frame = cap.read()
|
||||||
|
if not ret:
|
||||||
|
print("Failed to grab the initial frame")
|
||||||
|
cap.release()
|
||||||
|
|
||||||
|
# Let the user draw the initial bounding box
|
||||||
|
bbox = cv2.selectROI("Select ROI", frame, fromCenter=False, showCrosshair=True)
|
||||||
|
|
||||||
|
# Initialize Dlib correlation tracker
|
||||||
|
tracker = dlib.correlation_tracker()
|
||||||
|
tracker.start_track(frame, dlib.rectangle(bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]))
|
||||||
|
|
||||||
|
img_counter = 0 # To count frames and name images
|
||||||
|
frame_counter = 0 # To track the number of frames
|
||||||
|
start_time = time.time() # Start time for FPS calculation
|
||||||
|
|
||||||
|
try:
|
||||||
|
while True:
|
||||||
|
ret, frame = cap.read()
|
||||||
|
if not ret:
|
||||||
|
print("Failed to grab frame")
|
||||||
|
break
|
||||||
|
|
||||||
|
# Start measuring time per frame
|
||||||
|
frame_start_time = time.time()
|
||||||
|
|
||||||
|
# Update the tracker with the current frame
|
||||||
|
tracker.update(frame)
|
||||||
|
pos = tracker.get_position()
|
||||||
|
|
||||||
|
# Extract the bounding box from the tracker
|
||||||
|
x1, y1, x2, y2 = int(pos.left()), int(pos.top()), int(pos.right()), int(pos.bottom())
|
||||||
|
|
||||||
|
# Draw the updated tracking bounding box on the frame
|
||||||
|
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
||||||
|
|
||||||
|
# Show the frame with the tracking rectangle
|
||||||
|
cv2.imshow("Tracking", frame)
|
||||||
|
|
||||||
|
# Save every 10th frame
|
||||||
|
if frame_counter % 10 == 0:
|
||||||
|
img_name = f"tracked_frame_{img_counter}.jpg"
|
||||||
|
cv2.imwrite(img_name, frame)
|
||||||
|
print(f"Saved {img_name}")
|
||||||
|
img_counter += 1
|
||||||
|
|
||||||
|
# Calculate and print FPS
|
||||||
|
frame_time = time.time() - frame_start_time
|
||||||
|
fps = 1.0 / frame_time
|
||||||
|
print(f"FPS: {fps:.2f}")
|
||||||
|
|
||||||
|
frame_counter += 1
|
||||||
|
|
||||||
|
# Exit if 'q' is pressed
|
||||||
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
||||||
|
break
|
||||||
|
|
||||||
|
except KeyboardInterrupt:
|
||||||
|
print("Tracking interrupted by user")
|
||||||
|
|
||||||
|
cap.release()
|
||||||
|
cv2.destroyAllWindows()
|
||||||
|
|
||||||
|
# Calculate and print average FPS over all frames
|
||||||
|
end_time = time.time()
|
||||||
|
total_time = end_time - start_time
|
||||||
|
average_fps = frame_counter / total_time
|
||||||
|
print(f"Average FPS: {average_fps:.2f}")
|
||||||
|
|
||||||
|
print("Finished tracking")
|
||||||
|
|
@ -0,0 +1,116 @@
|
||||||
|
import cv2 as cv
|
||||||
|
import dlib
|
||||||
|
import numpy as np
|
||||||
|
import os
|
||||||
|
|
||||||
|
def match_feature_find_object(query_img, train_img, min_matches, results_file_name):
|
||||||
|
# Create an ORB object
|
||||||
|
orb = cv.ORB_create(nfeatures=100000)
|
||||||
|
|
||||||
|
features1, des1 = orb.detectAndCompute(query_img, None)
|
||||||
|
features2, des2 = orb.detectAndCompute(train_img, None)
|
||||||
|
|
||||||
|
# Create Brute-Force matcher object
|
||||||
|
bf = cv.BFMatcher(cv.NORM_HAMMING)
|
||||||
|
matches = bf.knnMatch(des1, des2, k=2)
|
||||||
|
|
||||||
|
# Nearest neighbour ratio test to find good matches
|
||||||
|
good = []
|
||||||
|
good_without_lists = []
|
||||||
|
matches = [match for match in matches if len(match) == 2]
|
||||||
|
for m, n in matches:
|
||||||
|
if m.distance < 0.8 * n.distance:
|
||||||
|
good.append([m])
|
||||||
|
good_without_lists.append(m)
|
||||||
|
|
||||||
|
if len(good) >= min_matches:
|
||||||
|
print(f"good: {len(good)}, {results_file_name}")
|
||||||
|
# Draw a polygon around the recognized object
|
||||||
|
src_pts = np.float32([features1[m.queryIdx].pt for m in good_without_lists]).reshape(-1, 1, 2)
|
||||||
|
dst_pts = np.float32([features2[m.trainIdx].pt for m in good_without_lists]).reshape(-1, 1, 2)
|
||||||
|
|
||||||
|
# Get the transformation matrix
|
||||||
|
M, _ = cv.findHomography(src_pts, dst_pts, cv.RANSAC, 5.0)
|
||||||
|
|
||||||
|
# Find the perspective transformation to get the corresponding points
|
||||||
|
h, w = query_img.shape[:2]
|
||||||
|
pts = np.float32([[0, 0], [0, h - 1], [w - 1, h - 1], [w - 1, 0]]).reshape(-1, 1, 2)
|
||||||
|
dst = cv.perspectiveTransform(pts, M)
|
||||||
|
|
||||||
|
train_img = cv.polylines(train_img, [np.int32(dst)], True, (0, 255, 0), 2, cv.LINE_AA)
|
||||||
|
|
||||||
|
# Return the bounding box of the detected object
|
||||||
|
x, y, w, h = cv.boundingRect(np.int32(dst))
|
||||||
|
return (x, y, w, h), train_img # Return bounding box and modified image
|
||||||
|
|
||||||
|
else:
|
||||||
|
print("Not enough good matches are found - {}/{}".format(len(good), min_matches))
|
||||||
|
return None, train_img # No detection, return None
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
input_folder = "extracted_frames/"
|
||||||
|
output_folder = "results_extracted_frames/"
|
||||||
|
new_object_size = (300, 300)
|
||||||
|
|
||||||
|
os.makedirs(input_folder, exist_ok=True)
|
||||||
|
os.makedirs(output_folder, exist_ok=True)
|
||||||
|
|
||||||
|
orb = cv.ORB_create()
|
||||||
|
query_img = cv.imread("img/drone.png")
|
||||||
|
query_img = cv.resize(query_img, new_object_size)
|
||||||
|
features = orb.detect(query_img, None)
|
||||||
|
f_img = cv.drawKeypoints(query_img, features, None, color=(0, 255, 0), flags=0)
|
||||||
|
# cv.imwrite(os.path.join(output_folder, "drone.png"), f_img) # Save the image with keypoints
|
||||||
|
|
||||||
|
# Initialize dlib correlation tracker
|
||||||
|
tracker = dlib.correlation_tracker()
|
||||||
|
is_tracking = False # Flag to indicate if we are tracking
|
||||||
|
|
||||||
|
# Sort the filenames to process them in sequence
|
||||||
|
filenames = sorted([f for f in os.listdir(input_folder) if f.endswith("png")])
|
||||||
|
|
||||||
|
# Previous bounding box (if the object is detected)
|
||||||
|
prev_bbox = None
|
||||||
|
|
||||||
|
for filename in filenames:
|
||||||
|
img = cv.imread(os.path.join(input_folder, filename))
|
||||||
|
|
||||||
|
if not is_tracking: # Try to detect using feature matching
|
||||||
|
bbox, modified_img = match_feature_find_object(query_img, img, 100, os.path.join(output_folder, filename))
|
||||||
|
|
||||||
|
if bbox is not None:
|
||||||
|
# Object detected, initialize the dlib tracker
|
||||||
|
x, y, w, h = bbox
|
||||||
|
tracker.start_track(img, dlib.rectangle(x, y, x + w, y + h))
|
||||||
|
is_tracking = True
|
||||||
|
prev_bbox = bbox
|
||||||
|
else:
|
||||||
|
# No detection, use the image as is
|
||||||
|
modified_img = img
|
||||||
|
|
||||||
|
else: # If we are tracking
|
||||||
|
# Update the tracker and get the new position
|
||||||
|
tracker.update(img)
|
||||||
|
pos = tracker.get_position()
|
||||||
|
x1, y1, x2, y2 = int(pos.left()), int(pos.top()), int(pos.right()), int(pos.bottom())
|
||||||
|
print(x1, y1, x2, y2)
|
||||||
|
|
||||||
|
# Draw the tracking bounding box
|
||||||
|
cv.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
||||||
|
|
||||||
|
# If the bounding box is valid, keep tracking
|
||||||
|
if (x2 - x1) > 0 and (y2 - y1) > 0:
|
||||||
|
prev_bbox = (x1, y1, x2 - x1, y2 - y1)
|
||||||
|
modified_img = img
|
||||||
|
else:
|
||||||
|
# If tracking fails, stop tracking and try detection again in the next frame
|
||||||
|
is_tracking = False
|
||||||
|
modified_img = img
|
||||||
|
|
||||||
|
# Save the result as a PNG image
|
||||||
|
# cv.imwrite(os.path.join(output_folder, filename), modified_img)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
|
|
@ -0,0 +1,95 @@
|
||||||
|
from PIL import Image, ImageOps
|
||||||
|
import random
|
||||||
|
import os
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
RANDOM_SEED = 42
|
||||||
|
random.seed(RANDOM_SEED)
|
||||||
|
|
||||||
|
def resize_image(image, target_size):
|
||||||
|
original_width, original_height = image.size
|
||||||
|
target_width, target_height = target_size
|
||||||
|
|
||||||
|
# Calculate the aspect ratios
|
||||||
|
width_ratio = target_width / original_width
|
||||||
|
height_ratio = target_height / original_height
|
||||||
|
|
||||||
|
# Use the smaller ratio to ensure the image fits within the target dimensions
|
||||||
|
scaling_factor = min(width_ratio, height_ratio)
|
||||||
|
|
||||||
|
# Calculate new dimensions
|
||||||
|
new_width = int(original_width * scaling_factor)
|
||||||
|
new_height = int(original_height * scaling_factor)
|
||||||
|
|
||||||
|
# Resize the image
|
||||||
|
return image.resize((new_width, new_height), Image.LANCZOS)
|
||||||
|
|
||||||
|
def rotate_image(image, angle):
|
||||||
|
return image.rotate(angle, expand=True, resample=Image.Resampling.BICUBIC)
|
||||||
|
|
||||||
|
def tilt_image(image, max_tilt_angle=10):
|
||||||
|
width, height = image.size
|
||||||
|
|
||||||
|
# Simulate tilting by applying an affine transformation
|
||||||
|
tilt_angle = random.uniform(-max_tilt_angle, max_tilt_angle)
|
||||||
|
x_shift = width * np.sin(np.radians(tilt_angle))
|
||||||
|
|
||||||
|
# Affine transformation matrix
|
||||||
|
if tilt_angle > 0:
|
||||||
|
transform_matrix = (1, x_shift / height, -x_shift / 2, 0, 1, 0)
|
||||||
|
else:
|
||||||
|
transform_matrix = (1, x_shift / height, x_shift / 2, 0, 1, 0)
|
||||||
|
|
||||||
|
return image.transform(image.size, Image.AFFINE, transform_matrix, resample=Image.Resampling.BICUBIC)
|
||||||
|
|
||||||
|
def generate_images(object_image, output_folder, background_image_path,
|
||||||
|
background_width=960, background_height=720, num_images=50,
|
||||||
|
max_rotation_angle=45, max_tilt_angle=10):
|
||||||
|
|
||||||
|
background_image = Image.open(background_image_path).convert("RGB")
|
||||||
|
background_image = background_image.resize((background_width, background_height))
|
||||||
|
|
||||||
|
for i in range(num_images):
|
||||||
|
# Create a copy of the background image
|
||||||
|
background = background_image.copy()
|
||||||
|
|
||||||
|
# Randomly rotate the object image
|
||||||
|
rotation_angle = random.uniform(-max_rotation_angle, max_rotation_angle)
|
||||||
|
rotated_image = rotate_image(object_image, rotation_angle)
|
||||||
|
|
||||||
|
# Apply random tilt to the object image
|
||||||
|
tilted_image = tilt_image(rotated_image, max_tilt_angle)
|
||||||
|
|
||||||
|
# Random position for the object
|
||||||
|
max_x = background_width - tilted_image.width
|
||||||
|
max_y = background_height - tilted_image.height
|
||||||
|
x = random.randint(0, max_x)
|
||||||
|
y = random.randint(0, max_y)
|
||||||
|
|
||||||
|
# Paste the object image onto the background
|
||||||
|
background.paste(tilted_image, (x, y), tilted_image)
|
||||||
|
|
||||||
|
# Save the new image
|
||||||
|
background.save(f"{output_folder}/variation_{i + 1}.png")
|
||||||
|
|
||||||
|
def main():
|
||||||
|
background_width = 960
|
||||||
|
background_height = 720
|
||||||
|
new_object_size = (300, 300)
|
||||||
|
|
||||||
|
os.makedirs("gen_data", exist_ok=True)
|
||||||
|
output_folder = "gen_data/"
|
||||||
|
|
||||||
|
# Load the object image (with transparency)
|
||||||
|
object_image = Image.open("img/drone.png").convert("RGBA")
|
||||||
|
object_image = resize_image(object_image, new_object_size)
|
||||||
|
|
||||||
|
# Load the background image from a second source
|
||||||
|
background_image_path = "img/background.jpg"
|
||||||
|
|
||||||
|
# Generate images with object placed on top of the background
|
||||||
|
generate_images(object_image, output_folder, background_image_path,
|
||||||
|
background_width, background_height, num_images=50)
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
|
|
@ -0,0 +1,4 @@
|
||||||
|
pillow==10.4.0
|
||||||
|
opencv-python==4.10.0.84
|
||||||
|
numpy==2.1.1
|
||||||
|
matplotlib==3.9.2
|
||||||
Loading…
Reference in New Issue